stubpy.imports

stubpy.imports

Import-statement analysis for .pyi header assembly.

Four 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

  4. collect_special_imports() — detect abc / dataclasses needs

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.

from module import * produces a single entry under the special key "*" mapping to the original from module import * statement. Callers that need to pass through star-imports verbatim can check for this key explicitly.

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'

>>> result = scan_import_statements("from demo import *")
>>> result["*"]
'from demo import *'
collect_typing_imports(content: str) → list[str][source]

Return sorted typing names actually referenced in content.

Dynamically scans typing.__all__ (computed once at module import time) rather than a fixed hard-coded list, so newly-added typing names are picked up automatically in future Python versions.

Uses whole-word matching (\b word boundaries) to avoid false positives from substrings (e.g. List does not match inside BlackList).

Also excludes names that are locally defined in the stub body (class names, function names, parameter names) to avoid incorrectly importing typing.Container when the stub defines class Container or importing typing.override when a method has a parameter named override.

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_typing_imports("class Container(Element): ...")
[]
collect_cross_imports(module: ModuleType, module_name: str, body: str, import_map: dict[str, str]) → list[str][source]

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

Scans the stub body for two patterns:

  1. Capitalised names used as base classes or in type annotations (e.g. Element in class Container(Element):).

  2. Dotted module references used in annotations (e.g. types in types.Length). The module name (lowercase) is looked up in import_map so that from demo import types is re-emitted when variables carry types.Length annotations.

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().

  • 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']

>>> body2 = "DEFAULT_WIDTH: types.Length"
>>> import_map2 = {"types": "from demo import types"}
>>> collect_cross_imports(m, "mymod", body2, import_map2)
['from demo import types']
collect_special_imports(body: str) → dict[str, list[str]][source]

Return a {module: [names]} dict of extra imports needed in body.

Checks for special constructs that require their own imports:

  • @abstractmethod or ABC base class → from abc import ...

  • @dataclass → from dataclasses import dataclass

Parameters:

body (str) – Generated stub body (class + function + variable stubs, no header).

Returns:

dict – Maps module name to a sorted list of names to import.

Examples

>>> collect_special_imports("@abstractmethod\ndef foo(): ...")
{'abc': ['abstractmethod']}
>>> collect_special_imports("@dataclass\nclass Foo: ...")
{'dataclasses': ['dataclass']}
>>> collect_special_imports("class Foo(ABC):\n    pass")
{'abc': ['ABC']}

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 inline ones inside function bodies — and works even when some imports fail to resolve at runtime.

Star imports (from module import *) are recorded under the reserved key "*" so callers can detect and handle them explicitly.

Typing candidates

collect_typing_imports() scans typing.__all__ at import time to build its candidate set. This means it automatically covers names added in future Python releases without any code changes. Matching uses whole-word boundaries so List is not falsely matched inside BlackList.

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.