wehub-resource-sync
394 lines
13 KiB
Python
394 lines
13 KiB
Python
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()
|