| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Ensure we perform lazy loading in vllm/__init__.py. |
| 4 | i.e: appears only within the `if typing.TYPE_CHECKING:` guard, |
| 5 | **except** for a short whitelist. |
| 6 | """ |
| 7 | |
| 8 | import ast |
| 9 | import sys |
| 10 | from collections.abc import Iterable |
| 11 | from pathlib import Path |
| 12 | from typing import Final |
| 13 | |
| 14 | INIT_PATH: Final = Path("vllm/__init__.py") |
| 15 | |
| 16 | # If you need to add items to whitelist, do it here. |
| 17 | ALLOWED_IMPORTS: Final[frozenset[str]] = frozenset( |
| 18 | { |
| 19 | "vllm.env_override", |
| 20 | } |
| 21 | ) |
| 22 | ALLOWED_FROM_MODULES: Final[frozenset[str]] = frozenset( |
| 23 | { |
| 24 | ".version", |
| 25 | } |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | def _is_internal(name: str | None, *, level: int = 0) -> bool: |
| 30 | if level > 0: |
| 31 | return True |
| 32 | if name is None: |
| 33 | return False |
| 34 | return name.startswith("vllm.") or name == "vllm" |
| 35 | |
| 36 | |
| 37 | def _fail(violations: Iterable[tuple[int, str]]) -> None: |
| 38 | print("ERROR: Disallowed eager imports in vllm/__init__.py:\n", file=sys.stderr) |
| 39 | for lineno, msg in violations: |
| 40 | print(f" Line {lineno}: {msg}", file=sys.stderr) |
| 41 | sys.exit(1) |
| 42 | |
| 43 | |
| 44 | def main() -> None: |
| 45 | source = INIT_PATH.read_text(encoding="utf-8") |
| 46 | tree = ast.parse(source, filename=str(INIT_PATH)) |
| 47 | |
| 48 | violations: list[tuple[int, str]] = [] |
| 49 | |
| 50 | class Visitor(ast.NodeVisitor): |
| 51 | def __init__(self) -> None: |
| 52 | super().__init__() |
| 53 | self._in_type_checking = False |
| 54 | |
| 55 | def visit_If(self, node: ast.If) -> None: |
| 56 | guard_is_type_checking = False |
| 57 | test = node.test |
| 58 | if isinstance(test, ast.Attribute) and isinstance(test.value, ast.Name): |
| 59 | guard_is_type_checking = ( |
| 60 | test.value.id == "typing" and test.attr == "TYPE_CHECKING" |
| 61 | ) |
| 62 | elif isinstance(test, ast.Name): |
| 63 | guard_is_type_checking = test.id == "TYPE_CHECKING" |
| 64 | |
| 65 | if guard_is_type_checking: |
| 66 | prev = self._in_type_checking |
| 67 | self._in_type_checking = True |
| 68 | for child in node.body: |
| 69 | self.visit(child) |
| 70 | self._in_type_checking = prev |
| 71 | for child in node.orelse: |
| 72 | self.visit(child) |
| 73 | else: |
| 74 | self.generic_visit(node) |
| 75 | |
| 76 | def visit_Import(self, node: ast.Import) -> None: |
| 77 | if self._in_type_checking: |
| 78 | return |
| 79 | for alias in node.names: |
| 80 | module_name = alias.name |
| 81 | if _is_internal(module_name) and module_name not in ALLOWED_IMPORTS: |
| 82 | violations.append( |
| 83 | ( |
| 84 | node.lineno, |
| 85 | f"import '{module_name}' must be inside typing.TYPE_CHECKING", # noqa: E501 |
| 86 | ) |
| 87 | ) |
| 88 | |
| 89 | def visit_ImportFrom(self, node: ast.ImportFrom) -> None: |
| 90 | if self._in_type_checking: |
| 91 | return |
| 92 | module_as_written = ("." * node.level) + (node.module or "") |
| 93 | if ( |
| 94 | _is_internal(node.module, level=node.level) |
| 95 | and module_as_written not in ALLOWED_FROM_MODULES |
| 96 | ): |
| 97 | violations.append( |
| 98 | ( |
| 99 | node.lineno, |
| 100 | f"from '{module_as_written}' import ... must be inside typing.TYPE_CHECKING", # noqa: E501 |
| 101 | ) |
| 102 | ) |
| 103 | |
| 104 | Visitor().visit(tree) |
| 105 | |
| 106 | if violations: |
| 107 | _fail(violations) |
| 108 | |
| 109 | |
| 110 | if __name__ == "__main__": |
| 111 | main() |