chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,393 @@
|
||||
import argparse
|
||||
import ast
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Set, Tuple
|
||||
|
||||
TRAIN_PACKAGES = set()
|
||||
|
||||
|
||||
def find_train_packages(base_train_dir: Path, patch_train_dir: Path) -> None:
|
||||
"""
|
||||
Find and initialize the global TRAIN_PACKAGES with all train package names from base_train_dir
|
||||
and patch_train_dir combined.
|
||||
"""
|
||||
global TRAIN_PACKAGES
|
||||
TRAIN_PACKAGES = set()
|
||||
|
||||
# Collect all packages under both base and patch train dirs
|
||||
package_files = list(base_train_dir.rglob("__init__.py")) + list(
|
||||
patch_train_dir.rglob("__init__.py")
|
||||
)
|
||||
base_dir = get_base_dir()
|
||||
for init_file in package_files:
|
||||
relative_path = init_file.relative_to(base_dir)
|
||||
dotted_module = str(relative_path.parent).replace("/", ".")
|
||||
TRAIN_PACKAGES.add(dotted_module)
|
||||
|
||||
|
||||
def is_train_package(module_str: str) -> bool:
|
||||
return module_str in TRAIN_PACKAGES
|
||||
|
||||
|
||||
def get_base_dir() -> Path:
|
||||
"""Return the filesystem path to the ray python directory."""
|
||||
current_file_path = Path(__file__).resolve()
|
||||
package_dir = current_file_path.parents[3]
|
||||
return package_dir
|
||||
|
||||
|
||||
def get_base_train_dir() -> Path:
|
||||
"""Return the filesystem path to the ray train directory."""
|
||||
return get_base_dir() / "ray/train"
|
||||
|
||||
|
||||
def does_overlap(main_module: str, module: str) -> bool:
|
||||
"""Checks if the init file of module is on the import path of main_module"""
|
||||
return main_module.startswith(f"{module}.") or main_module == module
|
||||
|
||||
|
||||
class Import:
|
||||
"""
|
||||
Represents an import statement.
|
||||
For example, 'from X import A, B' has module 'X' and names ['A', 'B'].
|
||||
Also supports 'import X'.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, module: str, names: List[str] = None, is_package: bool = False
|
||||
) -> None:
|
||||
self.is_package = is_package
|
||||
self.module = module
|
||||
self.names = names if names else []
|
||||
|
||||
|
||||
class ImportCollector(ast.NodeVisitor):
|
||||
"""
|
||||
An AST node visitor that collects all module-level imports from a Python source file.
|
||||
It traverses the AST and records module-level import statements (`import ...` and `from ... import ...`) that are not
|
||||
inside function or class definitions, and that are not guarded by `if TYPE_CHECKING` or `if typing.TYPE_CHECKING`
|
||||
blocks.
|
||||
"""
|
||||
|
||||
def __init__(self, module_name: str, is_package: bool) -> None:
|
||||
self._module_name = module_name
|
||||
self._is_package = is_package
|
||||
self.imports: Set[Import] = set()
|
||||
self.type_checking_imported = False
|
||||
|
||||
# --- private helpers ---
|
||||
|
||||
def _is_type_checking_test(self, expr: ast.AST) -> bool:
|
||||
"""Return True for `if TYPE_CHECKING` or `if typing.TYPE_CHECKING`."""
|
||||
|
||||
if (
|
||||
self.type_checking_imported
|
||||
and isinstance(expr, ast.Name)
|
||||
and expr.id == "TYPE_CHECKING"
|
||||
):
|
||||
return True
|
||||
elif (
|
||||
isinstance(expr, ast.Attribute)
|
||||
and isinstance(expr.value, ast.Name)
|
||||
and expr.value.id == "typing"
|
||||
and expr.attr == "TYPE_CHECKING"
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _get_package_parts(self) -> List[str]:
|
||||
parts = self._module_name.split(".")
|
||||
return parts if self._is_package else parts[:-1]
|
||||
|
||||
def _to_absolute_module(
|
||||
self, level: int, module_str: Optional[str]
|
||||
) -> Optional[str]:
|
||||
"""Construct the absolute module string from a relative import."""
|
||||
# Absolute import
|
||||
if level == 0:
|
||||
return module_str
|
||||
|
||||
package_parts = self._get_package_parts()
|
||||
|
||||
# If the relative import is out of bounds
|
||||
if level - 1 > len(package_parts):
|
||||
return None
|
||||
|
||||
# Base parts based on the level
|
||||
base_module_parts = (
|
||||
package_parts if level == 1 else package_parts[: -(level - 1)]
|
||||
)
|
||||
|
||||
# Construct absolute module string
|
||||
abs_module_parts = (
|
||||
base_module_parts + module_str.split(".")
|
||||
if module_str
|
||||
else base_module_parts
|
||||
)
|
||||
return ".".join(abs_module_parts)
|
||||
|
||||
# --- parsing functions ---
|
||||
|
||||
def visit_If(self, node: ast.If) -> None:
|
||||
# If the test is not TYPE_CHECKING, visit statement body
|
||||
if not self._is_type_checking_test(node.test):
|
||||
for stmt in node.body:
|
||||
self.visit(stmt)
|
||||
|
||||
# Also visit conditional branches
|
||||
for stmt in node.orelse:
|
||||
self.visit(stmt)
|
||||
|
||||
def visit_Import(self, node: ast.Import) -> None:
|
||||
for alias in node.names:
|
||||
if alias.name:
|
||||
self.imports.add(
|
||||
Import(module=alias.name, is_package=is_train_package(alias.name))
|
||||
)
|
||||
|
||||
def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
|
||||
import_str = self._to_absolute_module(node.level or 0, node.module)
|
||||
if not import_str:
|
||||
return
|
||||
|
||||
names = [alias.name for alias in node.names]
|
||||
self.imports.add(
|
||||
Import(
|
||||
module=import_str, is_package=is_train_package(import_str), names=names
|
||||
)
|
||||
)
|
||||
if "TYPE_CHECKING" in names and import_str == "typing":
|
||||
self.type_checking_imported = True
|
||||
|
||||
def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
|
||||
# Skip function contents
|
||||
return
|
||||
|
||||
def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None:
|
||||
# Skip function contents
|
||||
return
|
||||
|
||||
def visit_ClassDef(self, node: ast.ClassDef) -> None:
|
||||
# Skip class contents
|
||||
return
|
||||
|
||||
|
||||
def collect_imports(
|
||||
module_name: str, is_package: bool, source_text: str
|
||||
) -> Set[Import]:
|
||||
try:
|
||||
tree = ast.parse(source_text)
|
||||
except SyntaxError:
|
||||
print(f"Warning: Failed to parse {module_name} for circular imports")
|
||||
return set()
|
||||
collector = ImportCollector(module_name, is_package)
|
||||
collector.visit(tree)
|
||||
return collector.imports
|
||||
|
||||
|
||||
def to_module_name_and_is_package(py_file: Path) -> Tuple[str, bool]:
|
||||
"""
|
||||
Convert a Python file path to its corresponding module name and determine if it is a package.
|
||||
|
||||
Args:
|
||||
py_file: The path to the Python file.
|
||||
|
||||
Returns:
|
||||
Tuple[str, bool]: A tuple containing the module name as a string and a boolean indicating
|
||||
whether the module is a package (True if it is an __init__.py file).
|
||||
"""
|
||||
file_path = py_file.relative_to(get_base_dir())
|
||||
module_path = file_path.with_suffix("")
|
||||
module_parts = module_path.parts
|
||||
is_package = module_parts[-1] == "__init__"
|
||||
if is_package:
|
||||
module_parts = module_parts[:-1]
|
||||
module_str = ".".join(module_parts)
|
||||
return module_str, is_package
|
||||
|
||||
|
||||
def get_file_module_imports(
|
||||
files: List[Path], module_match_string: Optional[str] = None
|
||||
) -> Dict[str, List[Import]]:
|
||||
"""
|
||||
Collect and return the module-level imports for a list of Python files.
|
||||
|
||||
Args:
|
||||
files: A list of Path objects representing Python files to analyze.
|
||||
module_match_string: An optional string to filter imports. Only imports
|
||||
containing this string will be included in the result.
|
||||
|
||||
Returns:
|
||||
A dictionary mapping module names to a list of their import statements.
|
||||
The module names are derived from the file paths, and the import statements
|
||||
are filtered based on the optional module_match_string.
|
||||
"""
|
||||
module_imports: Dict[str, List[Import]] = {}
|
||||
|
||||
# Collect the imports for each python file
|
||||
for py_file in files:
|
||||
try:
|
||||
module_name, is_package = to_module_name_and_is_package(py_file)
|
||||
src = py_file.read_text(encoding="utf-8", errors="ignore")
|
||||
imports = collect_imports(module_name, is_package, src)
|
||||
module_imports[module_name] = [
|
||||
stmt
|
||||
for stmt in imports
|
||||
if module_match_string is None or module_match_string in stmt.module
|
||||
]
|
||||
except Exception:
|
||||
continue
|
||||
return module_imports
|
||||
|
||||
|
||||
def convert_to_file_paths(imports: List[Import]) -> List[Path]:
|
||||
"""
|
||||
Convert a list of import strings to a list of file paths.
|
||||
|
||||
Args:
|
||||
imports: A list of Import objects
|
||||
|
||||
Returns:
|
||||
A list of file paths.
|
||||
"""
|
||||
base_dir = get_base_dir()
|
||||
file_paths = []
|
||||
for imp in imports:
|
||||
if imp.is_package:
|
||||
relative_path = imp.module.replace(".", "/") + "/__init__.py"
|
||||
else:
|
||||
relative_path = imp.module.replace(".", "/") + ".py"
|
||||
file_paths.append(base_dir / relative_path)
|
||||
return file_paths
|
||||
|
||||
|
||||
def expand_to_include_reexports(import_map: Dict[str, List[Import]]) -> None:
|
||||
"""
|
||||
Expands the set of imports for a given import map to include the modules resulting from reexports.
|
||||
So if in the base train module, there is "from x import a, b" and x is a package, then this function
|
||||
will explore the __init__.py of x and include the modules a and b were reexported from in the import map.
|
||||
"""
|
||||
for module, base_imports in import_map.items():
|
||||
# Get only the package imports
|
||||
packages = [imp for imp in base_imports if imp.is_package]
|
||||
package_files = convert_to_file_paths(packages)
|
||||
reexports = get_file_module_imports(package_files)
|
||||
|
||||
agg_reexports = []
|
||||
# Filter patch init file imports to those that only contain the right names
|
||||
for base_import in base_imports:
|
||||
if base_import.module in reexports:
|
||||
import_list = reexports[base_import.module]
|
||||
target_reexports = [
|
||||
imp
|
||||
for imp in import_list
|
||||
if set(imp.names) & set(base_import.names)
|
||||
]
|
||||
agg_reexports.extend(target_reexports)
|
||||
|
||||
# Expand modules to include reexported modules
|
||||
import_map[module].extend(agg_reexports)
|
||||
|
||||
|
||||
def check_violations(
|
||||
base_train_patching_imports: Dict[str, List[Import]], patch_dir: Path
|
||||
) -> List[str]:
|
||||
"""
|
||||
Check for circular import violations between base and patch train modules.
|
||||
|
||||
Args:
|
||||
base_train_patching_imports: A dictionary mapping base train module names to their imports.
|
||||
patch_dir: The directory path containing patch train modules.
|
||||
|
||||
Returns:
|
||||
A list of strings describing any circular import violations found.
|
||||
"""
|
||||
violations: List[str] = []
|
||||
|
||||
# Get the imports from the patch train init files
|
||||
patch_train_init_files = list(patch_dir.rglob("__init__.py"))
|
||||
patch_train_init_imports = get_file_module_imports(
|
||||
patch_train_init_files, module_match_string="ray.train"
|
||||
)
|
||||
|
||||
# Expand the imports to include reexports
|
||||
expand_to_include_reexports(base_train_patching_imports)
|
||||
|
||||
# Process each patch train init module for violations
|
||||
for base_train_init_module, imports in base_train_patching_imports.items():
|
||||
|
||||
# Get the imports from the patch train files
|
||||
patch_train_files = convert_to_file_paths(imports)
|
||||
patch_train_file_imports = get_file_module_imports(
|
||||
patch_train_files, module_match_string="ray.train"
|
||||
)
|
||||
|
||||
for patch_module, imports in patch_train_file_imports.items():
|
||||
# Skip if the base train init module is in the import path of the patch module
|
||||
if does_overlap(patch_module, base_train_init_module):
|
||||
continue
|
||||
|
||||
# Skip if the patch train module init file imports the base train init module
|
||||
patch_init_module = (
|
||||
".".join(patch_module.split(".")[:-1])
|
||||
if not is_train_package(patch_module)
|
||||
else patch_module
|
||||
)
|
||||
patch_init_imports = patch_train_init_imports.get(patch_init_module, [])
|
||||
if any(
|
||||
does_overlap(imp.module, base_train_init_module)
|
||||
for imp in patch_init_imports
|
||||
):
|
||||
continue
|
||||
|
||||
for patch_import in imports:
|
||||
# If any of those v1 imports go through the init file, then it is a violation
|
||||
if does_overlap(patch_import.module, base_train_init_module):
|
||||
violations.append(
|
||||
f"circular-import-train: Circular import between {base_train_init_module} (importing {patch_module}) and {patch_module} (importing {patch_import.module}). Resolve by importing {base_train_init_module} in the __init__.py of {patch_init_module}."
|
||||
)
|
||||
|
||||
return violations
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--patch_dir",
|
||||
default="ray/train/v2",
|
||||
help="Path to the directory containing patching contents",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Get train directory paths
|
||||
base_dir = get_base_dir()
|
||||
base_train_dir = get_base_train_dir()
|
||||
patch_train_dir = base_dir / Path(args.patch_dir)
|
||||
|
||||
# Find and save all train packages in global TRAIN_PACKAGES for reference
|
||||
find_train_packages(base_train_dir, patch_train_dir)
|
||||
|
||||
# Collect all base train init files
|
||||
base_train_init_files = [
|
||||
f
|
||||
for f in base_train_dir.rglob("__init__.py")
|
||||
if not f.is_relative_to(patch_train_dir)
|
||||
]
|
||||
|
||||
# Get the patching imports in the base train init files
|
||||
dotted_module_prefix = str(patch_train_dir.relative_to(base_dir)).replace("/", ".")
|
||||
patching_imports = get_file_module_imports(
|
||||
base_train_init_files, module_match_string=dotted_module_prefix
|
||||
)
|
||||
|
||||
# Collect all violations based off the patching imports
|
||||
violations = check_violations(patching_imports, patch_train_dir)
|
||||
if violations:
|
||||
print("\n".join(violations))
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user