feat: enhance project context and file handling in linter

Author: AlbertoCentonze

.gitignore

@@ -14,3 +14,5 @@ site
 .ruff_cache
 .devcontainer
 .claude
+.mypy_cache
+.pytest_cache

README.md

@@ -10,8 +10,6 @@ You can install Natrix using pip:
 pip install natrix
 ```
 
-Currently natrix requires [`uv`](https://docs.astral.sh/uv/) to be installed to function correctly.
-
 ### Vyper Version Compatibility
 
 Natrix supports Vyper compiler versions 0.4.0 and above.

docs/api/index.md

@@ -147,6 +147,29 @@ for access in func_node.memory_accesses:
     print(f"{access.type}: {access.var}")  # "read: balance"
 ```
 
+## Project Context API
+
+The `ProjectContext` class manages the dependency graph and compilation of all modules in a project:
+
+```python
+from pathlib import Path
+from natrix.context import ProjectContext
+
+# Create context with all project files
+vy_files = [Path("contract1.vy"), Path("contract2.vy")]
+context = ProjectContext(vy_files, extra_paths=(Path("lib"),))
+
+# Get module information
+module = context.get_module(Path("contract1.vy"))
+if module:
+    print(f"Dependencies: {module.dependencies}")
+    print(f"Dependents: {module.dependents}")
+
+# Get dependency relationships
+deps = context.get_dependencies_of(Path("contract1.vy"))
+dependents = context.get_dependents_of(Path("contract1.vy"))
+```
+
 ## Rule Development API
 
 ### Base Rule Class
@@ -226,15 +249,17 @@ class MyRule(BaseRule):
 Utilities for working with Vyper compilation:
 
 ```python
+from pathlib import Path
 from natrix.ast_tools import parse_file, vyper_compile, VyperASTVisitor
+from natrix.ast_node import Node
 
 # Parse a Vyper file to AST
-ast_data = parse_file("contract.vy")
-root_node = Node(ast_data)
+ast_data = parse_file(Path("contract.vy"))
+root_node = Node(ast_data["ast"])
 
 # Compile with specific format
-ast_only = vyper_compile("contract.vy", "annotated_ast")
-metadata = vyper_compile("contract.vy", "metadata")
+ast_only = vyper_compile(Path("contract.vy"), "annotated_ast")
+metadata = vyper_compile(Path("contract.vy"), "metadata")
 ```
 
 ### VyperASTVisitor
@@ -321,6 +346,8 @@ class FunctionNamingRule(BaseRule):
 
 ```python
 import pytest
+from pathlib import Path
+from natrix.context import ProjectContext
 from natrix.ast_tools import parse_file
 from natrix.ast_node import Node
 
@@ -337,12 +364,14 @@ def good_function_name():  # Should pass
 """
 
     # Parse and test
-    with open("test.vy", "w") as f:
+    test_path = Path("test.vy")
+    with open(test_path, "w") as f:
         f.write(test_contract)
 
-    ast_data = parse_file("test.vy")
+    # Create context and run rule
+    context = ProjectContext([test_path])
     rule = FunctionNamingRule()
-    issues = rule.run(ast_data)
+    issues = rule.run(context, test_path)
 
     assert len(issues) == 1
     assert "badFunctionName" in issues[0].message

natrix/__init__.py

@@ -16,10 +16,13 @@
     import tomli as tomllib
 
 from natrix.__version__ import __version__
-from natrix.ast_tools import parse_file
 from natrix.codegen import generate_exports
+from natrix.context import ProjectContext
 from natrix.rules.common import Issue, RuleRegistry
 
+# Vyper file extensions
+VYPER_EXTENSIONS = (".vy", ".vyi")
+
 
 class OutputFormatter:
     """Handles output formatting for both CLI and JSON modes."""
@@ -67,14 +70,12 @@ def print_summary(self, has_issues: bool) -> None:
 
 
 def lint_file(
-    file_path: str,
+    file_path: Path,
+    project_context: ProjectContext,
     formatter: OutputFormatter,
     disabled_rules: set[str] | None = None,
-    extra_paths: tuple[str, ...] = (),
 ) -> list[Issue]:
-    """Lint a single Vyper file with the given rules configuration."""
-    ast = parse_file(file_path, extra_paths=extra_paths)
-
+    """Lint a single Vyper file using the pre-built project context."""
     if disabled_rules is None:
         disabled_rules = set()
 
@@ -85,7 +86,7 @@ def lint_file(
     issues = []
     for rule in rules:
         try:
-            rule_issues = rule.run(ast)
+            rule_issues = rule.run(project_context, file_path)
             issues.extend(
                 [issue for issue in rule_issues if issue.code not in disabled_rules]
             )
@@ -99,30 +100,31 @@ def lint_file(
     return issues
 
 
-def find_vy_files(directory: str) -> list[str]:
-    # Recursively find all .vy files in the given directory,
+def find_vy_files(directory: Path) -> list[Path]:
+    # Recursively find all Vyper files (.vy and .vyi) in the given directory,
     # excluding specified directories
     vy_files = []
     for root, _, files in os.walk(directory):
-        # Collect all .vy files
+        # Collect all Vyper files
         for file in files:
-            if file.endswith(".vy"):
-                vy_files.append(str(Path(root) / file))
+            file_path = Path(root) / file
+            if file_path.suffix in VYPER_EXTENSIONS:
+                vy_files.append(file_path)
 
     return vy_files
 
 
-def get_project_root() -> str:
+def get_project_root() -> Path:
     """Get the project root directory, which contains the pyproject.toml file."""
     # First try to find it from the current working directory upwards
     current_path = Path.cwd().resolve()
 
     for parent in [current_path, *current_path.parents]:
         if (parent / "pyproject.toml").exists():
-            return str(parent)
+            return parent
 
     # If not found, default to the current directory
-    return str(Path.cwd())
+    return Path.cwd()
 
 
 def read_pyproject_config() -> dict[str, Any]:
@@ -137,7 +139,7 @@ def read_pyproject_config() -> dict[str, Any]:
     try:
         # Find the project root directory
         project_root = get_project_root()
-        pyproject_path = Path(project_root) / "pyproject.toml"
+        pyproject_path = project_root / "pyproject.toml"
 
         if pyproject_path.exists():
             with pyproject_path.open("rb") as f:
@@ -149,7 +151,7 @@ def read_pyproject_config() -> dict[str, Any]:
                     ):
                         # Make paths relative to pyproject.toml location
                         config["files"] = [
-                            str((Path(project_root) / path).resolve())
+                            (project_root / path).resolve()
                             for path in natrix_config["files"]
                         ]
                     if "disabled_rules" in natrix_config and isinstance(
@@ -167,7 +169,7 @@ def read_pyproject_config() -> dict[str, Any]:
                     ):
                         # Make paths relative to pyproject.toml location
                         config["path"] = [
-                            str((Path(project_root) / path).resolve())
+                            (project_root / path).resolve()
                             for path in natrix_config["path"]
                         ]
     except Exception as e:
@@ -280,9 +282,9 @@ def main() -> None:
     if args.command == "codegen":
         if args.codegen_command == "exports":
             # Get extra paths if provided
-            extra_paths = tuple(args.path) if args.path else ()
+            extra_paths = tuple(Path(p) for p in args.path) if args.path else ()
             # Generate and print exports
-            exports = generate_exports(args.file_path, extra_paths)
+            exports = generate_exports(Path(args.file_path), extra_paths)
             print(exports)
             sys.exit(0)
         else:
@@ -361,36 +363,43 @@ def main() -> None:
     if args.files:
         all_vy_files = []
         for path in args.files:
-            path_obj = Path(path)
-            if path_obj.is_file() and path.endswith(".vy"):
+            path = Path(path)
+            if path.is_file() and path.suffix in VYPER_EXTENSIONS:
                 all_vy_files.append(path)
-            elif path_obj.is_dir():
+            elif path.is_dir():
                 dir_vy_files = find_vy_files(path)
                 if not dir_vy_files:
-                    formatter.print(f"No .vy files found in the directory: {path}")
+                    formatter.print(f"No Vyper files found in the directory: {path}")
                 all_vy_files.extend(dir_vy_files)
             else:
                 formatter.print(
-                    f"Provided path is not a valid .vy file or directory: {path}"
+                    f"Provided path is not a valid Vyper file or directory: {path}"
                 )
 
         if not all_vy_files:
-            formatter.print("No valid .vy files to lint.")
+            formatter.print("No valid Vyper files to lint.")
             sys.exit(1)
     else:
-        # If no paths are provided, search for .vy files in the current
+        # If no paths are provided, search for Vyper files in the current
         # directory recursively
-        all_vy_files = find_vy_files(".")
+        all_vy_files = find_vy_files(Path())
 
         if not all_vy_files:
-            formatter.print("No .vy files found in the current directory.")
+            formatter.print("No Vyper files found in the current directory.")
             sys.exit(1)
 
+    # Create ProjectContext with all files
+    formatter.print("Building project dependency graph...")
+    project_context = ProjectContext(
+        all_vy_files,
+        extra_paths=tuple(Path(p) if isinstance(p, str) else p for p in extra_paths),
+    )
+
     # Collect all issues from all files
     all_issues: list[Issue] = []
     for file in all_vy_files:
         file_issues = lint_file(
-            file, formatter, disabled_rules, extra_paths=extra_paths
+            Path(file).resolve(), project_context, formatter, disabled_rules
         )
         all_issues.extend(file_issues)
 

natrix/ast_tools.py

@@ -38,7 +38,7 @@ def _check_vyper_version() -> None:
         ) from e
 
 
-def _obtain_sys_path() -> list[str]:
+def _obtain_sys_path() -> list[Path]:
     """
     Obtain all the system paths in which the compiler would
     normally look for modules. This allows to lint vyper files
@@ -55,45 +55,37 @@ def _obtain_sys_path() -> list[str]:
     # Remove the first element of the list which is an empty string.
     paths = ast.literal_eval(stdout)[1:]
 
-    # Remove paths that do not exist as the vyper compiler will throw an error.
-    valid_paths = [path for path in paths if Path(path).exists()]
+    return [Path(path) for path in paths]
 
-    return valid_paths
 
-
-def _obtain_default_paths() -> list[str]:
+def _obtain_default_paths() -> list[Path]:
     """
     Obtain default paths for Vyper imports.
-    Only returns paths that actually exist on the system.
     """
     # List of default paths to check
     default_paths = [
-        "lib/pypi",  # Default dependency folder for moccasin
+        Path("lib/pypi"),  # Default dependency folder for moccasin
         # Add more paths here in the future
     ]
 
-    # Return only existing paths
-    existing_paths = []
-    for path in default_paths:
-        path_obj = Path(path)
-        if path_obj.exists() and path_obj.is_dir():
-            existing_paths.append(path)
-
-    return existing_paths
+    return default_paths
 
 
 def vyper_compile(
-    filename: str, formatting: str, extra_paths: tuple[str, ...] = ()
+    filename: Path, formatting: str, extra_paths: tuple[Path, ...] = ()
 ) -> dict[str, Any] | list[dict[str, Any]]:
     _check_vyper_version()
 
     # Combine all paths
     all_paths = _obtain_sys_path() + _obtain_default_paths() + list(extra_paths)
 
-    # Convert to compiler flags
-    path_flags = [item for p in all_paths for item in ["-p", p]]
+    # Filter out non-existent paths as the vyper compiler will throw an error
+    valid_paths = [p for p in all_paths if p.exists()]
 
-    command = ["vyper", "-f", formatting, filename, *path_flags]
+    # Convert to compiler flags (vyper expects strings)
+    path_flags = [item for p in valid_paths for item in ["-p", str(p)]]
+
+    command = ["vyper", "-f", formatting, str(filename), *path_flags]
 
     process = subprocess.Popen(
         command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
@@ -113,11 +105,15 @@ def vyper_compile(
         ) from e
 
 
-def parse_file(file_path: str, extra_paths: tuple[str, ...] = ()) -> dict[str, Any]:
+def parse_file(file_path: Path, extra_paths: tuple[Path, ...] = ()) -> dict[str, Any]:
     ast = vyper_compile(file_path, "annotated_ast", extra_paths=extra_paths)
     # For annotated_ast, vyper_compile returns a dict
     assert isinstance(ast, dict)
 
+    # For interface files (.vyi), we only compile to AST, not metadata
+    if file_path.suffix == ".vyi":
+        return ast
+
     # Try to compile metadata, but handle InitializerException gracefully
     # This happens when a module uses deferred initialization (uses: module_name)
     try:
@@ -160,7 +156,7 @@ def parse_source(source_code: str) -> dict[str, Any]:
 
     try:
         # Parse the temporary file
-        result = parse_file(temp_file_path)
+        result = parse_file(Path(temp_file_path))
         return result
     finally:
         # Clean up the temporary file

natrix/codegen.py

@@ -5,7 +5,7 @@
 from natrix.ast_tools import vyper_compile
 
 
-def generate_exports(file_path: str, extra_paths: tuple[str, ...]) -> str:
+def generate_exports(file_path: Path, extra_paths: tuple[Path, ...]) -> str:
     """Generate explicit exports for a Vyper contract.
 
     Args:
@@ -16,7 +16,7 @@ def generate_exports(file_path: str, extra_paths: tuple[str, ...]) -> str:
         A string containing the exports declaration
     """
     # Extract module name from file path
-    module_name = Path(file_path).stem
+    module_name = file_path.stem
 
     # Get the ABI from vyper
     abi = vyper_compile(file_path, "abi", extra_paths=extra_paths)