stubpy.imports

stubpy.imports

Import-statement analysis for .pyi header assembly.

Three functions cover the full import pipeline:

  1. scan_import_statements() — parse source AST → {name: stmt}

  2. collect_typing_imports() — find used typing names in stub body

  3. collect_cross_imports() — find cross-file class imports to re-emit

scan_import_statements(source: str) → dict[str, str][source]

Parse source and return a {local_name: import_statement} map.

Walks the AST of source and builds one entry per imported name. Both from … import … and plain import … forms are supported, including as aliases. A single from statement importing multiple names produces one entry per name.

Parameters:

source (str) – Raw Python source text to parse.

Returns:

dict – Maps each locally-bound name to its minimal import statement. Returns an empty dict on SyntaxError.

Examples

>>> result = scan_import_statements("from demo import types\nimport os")
>>> result["types"]
'from demo import types'
>>> result["os"]
'import os'

>>> result = scan_import_statements("from typing import Optional as Opt")
>>> result["Opt"]
'from typing import Optional as Opt'
collect_typing_imports(content: str) → list[str][source]

Return sorted typing names actually referenced in content.

Checks each candidate name against content so only used names appear in the from typing import … header line of the generated stub.

Parameters:

content (str) – Generated stub body text (class and method stubs, without the header lines).

Returns:

list of str – Sorted list of typing names present in content. Empty list if none are used.

Examples

>>> collect_typing_imports("def foo(x: Optional[str]) -> None: ...")
['Optional']

>>> collect_typing_imports("x: List[int], y: Dict[str, Any]")
['Any', 'Dict', 'List']

>>> collect_typing_imports("def foo(x: int) -> None: ...")
[]
collect_cross_imports(module: ModuleType, module_name: str, body: str, import_map: dict[str, str]) → list[str][source]

Find cross-file class imports that must be re-emitted in the .pyi header.

Scans the stub body for names used as base classes or type annotations, then resolves each against import_map to find statements that come from other local modules.

A name is included when all of the following hold:

  • It appears as a base class (class Foo(Bar)) or as a capitalised annotation name after ": " or "-> ".

  • Its import statement exists in import_map.

  • The statement does not start with a stdlib/typing prefix.

  • The name is not defined in the module being stubbed.

Parameters:

  • module (types.ModuleType) – The loaded source module, used to test whether a name is local.

  • module_name (str) – Synthetic module name from load_module(). Names whose __module__ matches this are treated as local.

  • body (str) – Already-generated stub class bodies (without the header lines).

  • import_map (dict) – Result of scan_import_statements() for the source file.

Returns:

list of str – Deduplicated import statement strings safe to add verbatim to the .pyi header. Empty list if nothing needs importing.

Examples

>>> import types as _t
>>> m = _t.ModuleType("mymod")
>>> body = "class Container(Element):\n    pass"
>>> import_map = {"Element": "from demo.element import Element"}
>>> collect_cross_imports(m, "mymod", body, import_map)
['from demo.element import Element']

Import scanning

scan_import_statements() uses ast.parse() on the raw source text rather than inspecting the loaded module. This means it captures all imports — including conditional ones — and works even when some imports fail to resolve at runtime.

The candidates checked by collect_typing_imports() are: Any, Callable, ClassVar, Dict, FrozenSet, Iterator, List, Literal, Optional, Sequence, Set, Tuple, Type, Union.

Cross-import heuristic

collect_cross_imports() uses two regex patterns on the stub body:

  1. Base classes — class \\w+\\(([^)]+)\\)

  2. Annotation names — capitalised words after ": " or "-> "

Matches are filtered against the import map, stdlib prefixes, and names defined in the module being stubbed.