shpe 0.2.0__py3-none-any.whl

shpe/__init__.py

@@ -0,0 +1,5 @@
+from .checker import Checker
+from .diagnostics import ErrorCode
+from .environment import ShapeState
+
+__all__ = ["Checker", "ErrorCode", "ShapeState"]

shpe/__main__.py

@@ -0,0 +1,4 @@
+from shpe.cli import main
+
+if __name__ == "__main__":
+    main()

shpe/checker.py

@@ -0,0 +1,159 @@
+import ast
+from typing import Any
+
+from shpe.diagnostics import Diagnostics, ErrorCode
+from shpe.environment import Environment, ShapeState
+from shpe.extractor import Extractor
+from shpe.resolver import Resolver
+
+
+class Checker(ast.NodeVisitor):
+    """AST visitor that walks through Python code to infer and check NumPy shapes"""
+
+    def __init__(self) -> None:
+        self.env = Environment()
+        self.diagnostics = Diagnostics()
+        self.extractor = Extractor(self.env, self.diagnostics)
+        self.resolver = Resolver(
+            self.extractor,
+            self.env,
+            self.diagnostics,
+            self,
+        )
+
+    @property
+    def shapes(self) -> dict[str, tuple[Any, ...] | ShapeState]:
+        """Convenience property for tests and CLI to access global shapes."""
+        return self.env.shapes
+
+    @property
+    def scalar_values(self) -> dict[str, int | float | ShapeState]:
+        """Convenience property for tests and CLI to access global scalar values."""
+        return self.env.scalar_values
+
+    @property
+    def errors(self) -> list[dict[str, Any]]:
+        """Convenience property for tests and CLI to access collected errors."""
+        return self.diagnostics.errors
+
+    @property
+    def warnings(self) -> list[dict[str, Any]]:
+        """Convenience property for tests and CLI to access collected warnings."""
+        return self.diagnostics.warnings
+
+    @property
+    def inlay_hints(self) -> list[dict[str, Any]]:
+        """Convenience property for tests and CLI to access collected inlay hints."""
+        return self.diagnostics.inlay_hints
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        """Handles explicit annotated variable assignments."""
+
+        if not isinstance(node.target, ast.Name):
+            return
+
+        var_name = node.target.id
+
+        if isinstance(node.value, ast.Constant) and isinstance(
+            node.value.value, (int, float)
+        ):
+            self.env.set_scalar(var_name, node.value.value)
+
+        annotated_shape = self.extractor.annotation_shape(node.annotation)
+        inferred_shape = self.resolver.shape(node.value) if node.value else None
+
+        resolved_shape = (
+            annotated_shape if annotated_shape is not None else inferred_shape
+        )
+        if resolved_shape is None:
+            resolved_shape = ShapeState.UNKNOWN
+        self.env.set_shape(var_name, resolved_shape)
+
+        if inferred_shape is not None and inferred_shape is not ShapeState.UNKNOWN:
+            self.diagnostics.hint(node, inferred_shape)
+        if (
+            annotated_shape is not None
+            and inferred_shape is not None
+            and annotated_shape is not ShapeState.UNKNOWN
+            and inferred_shape is not ShapeState.UNKNOWN
+            and annotated_shape != inferred_shape
+        ):
+            self.diagnostics.error(
+                node,
+                ErrorCode.ANNOTATION,
+                (
+                    f"{var_name} annotated as {annotated_shape}, "
+                    f"but expression has the shape {inferred_shape}. "
+                ),
+            )
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        """Handles implicit variable assignments."""
+
+        if (
+            node.value
+            and isinstance(node.value, ast.Constant)
+            and isinstance(node.value.value, (int, float))
+        ):
+            for target in node.targets:
+                if isinstance(target, ast.Name):
+                    self.env.set_scalar(target.id, node.value.value)
+
+        inferred_shape = self.resolver.shape(node.value) if node.value else None
+
+        if inferred_shape is not None and inferred_shape is not ShapeState.UNKNOWN:
+            self.diagnostics.hint(node, inferred_shape)
+
+        shape_to_set = (
+            inferred_shape if inferred_shape is not None else ShapeState.UNKNOWN
+        )
+        for target in node.targets:
+            if isinstance(target, ast.Name):
+                self.env.set_shape(target.id, shape_to_set)
+
+    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        """Tracks custom function definitions."""
+        self.env.set_function(node.name, node)
+        self.env.push_child()
+        self.resolver.user_function_def(node)
+        self.env.pop_child()
+
+    def visit_If(self, node: ast.If) -> None:
+        """Handles if statements by tainting variables modified inside branches."""
+        test_str = ast.unparse(node.test).replace('"', "'")
+
+        # Evaluate normally for the standard main block
+        if test_str == "__name__ == '__main__'":
+            self.generic_visit(node)
+        else:
+            self._visit_tainted_block(node.body)
+            if node.orelse:
+                self._visit_tainted_block(node.orelse)
+
+    def visit_For(self, node: ast.For) -> None:
+        """Handles for loops by tainting variables modified inside the loop."""
+        self._visit_tainted_block(node.body)
+
+    def visit_While(self, node: ast.While) -> None:
+        """Handles while loops by tainting variables modified inside the loop."""
+        self._visit_tainted_block(node.body)
+
+    def visit_AsyncFor(self, node: ast.AsyncFor) -> None:
+        """Handles async for loops."""
+        self._visit_tainted_block(node.body)
+
+    def _visit_tainted_block(self, body: list[ast.stmt]) -> None:
+        """Walks a block of code and forces any assigned variables to Unknown.
+        Used for complex code blocks where we cannot guarantee shape inference.
+        We would rather have Unkown than incorrect shapes.
+
+        Args:
+            body: A list of AST statements representing the code block.
+        """
+        for stmt in body:
+            self.visit(stmt)
+            for target in ast.walk(stmt):
+                if isinstance(target, ast.Name) and isinstance(target.ctx, ast.Store):
+                    self.env.set_shape(target.id, ShapeState.UNKNOWN)
+                    if target.id in self.env.scalar_values:
+                        self.env.scalar_values[target.id] = ShapeState.UNKNOWN

shpe/cli.py

@@ -0,0 +1,193 @@
+import argparse
+import ast
+import importlib.metadata
+import sys
+from pathlib import Path
+from typing import Any
+
+from shpe.checker import Checker
+from shpe.environment import ShapeState
+
+try:
+    __version__ = importlib.metadata.version("shpe")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "unknown"
+
+
+def main() -> None:
+    """CLI entry point for the local Python shape checker."""
+    parser = argparse.ArgumentParser(description="A local Python shape checker.")
+    parser.add_argument(
+        "paths", nargs="+", type=Path, help="Files or directories to check"
+    )
+    parser.add_argument(
+        "--version", action="version", version=f"%(prog)s {__version__}"
+    )
+    parser.add_argument(
+        "--show-shapes",
+        action="store_true",
+        help="Display inferred NumPy shapes for all variables",
+    )
+    args = parser.parse_args()
+
+    files = discover_files(args.paths)
+    all_errors: list[str] = []
+    all_shapes: dict[Path, dict[str, tuple[Any, ...] | ShapeState]] = {}
+    all_scalars: dict[Path, dict[str, int | float | ShapeState]] = {}
+
+    for filepath in files:
+        file_errors, shapes, scalars = _process_file(filepath)
+        all_errors.extend(file_errors)
+        if shapes:
+            all_shapes[filepath] = shapes
+        if scalars:
+            all_scalars[filepath] = scalars
+
+    if args.show_shapes and all_shapes:
+        _print_shape_report(all_shapes, all_scalars)
+
+    for err in all_errors:
+        print(err)
+
+    if all_errors:
+        print(f"\nFound {len(all_errors)} error(s) across {len(files)} file(s).")
+        sys.exit(1)
+    else:
+        print(f"Success: Checked {len(files)} file(s), no shape errors found.")
+        sys.exit(0)
+
+
+def discover_files(paths: list[Path]) -> list[Path]:
+    """Discovers and collects all Python (.py) source files from the provided paths.
+
+    Args:
+        paths: A list of file or directory Path objects to scan.
+
+    Returns:
+        A list of resolved Path objects pointing to valid Python source files.
+    """
+    files_to_check = []
+    for path in paths:
+        if not path.exists():
+            print(f"Error: Path '{path}' not found.", file=sys.stderr)
+            continue
+
+        if path.is_file() and path.suffix == ".py":
+            files_to_check.append(path)
+        elif path.is_dir():
+            for subpath in path.rglob("*.py"):
+                # Skip hidden directories like .git, .venv, __pycache__
+                if not any(part.startswith(".") for part in subpath.parts):
+                    files_to_check.append(subpath)
+
+    return files_to_check
+
+
+def _process_file(
+    filepath: Path,
+) -> tuple[
+    list[str],
+    dict[str, tuple[Any, ...] | ShapeState],
+    dict[str, int | float | ShapeState],
+]:
+    """Processes a single source file, running syntax checks and shape analysis.
+
+    Args:
+        filepath: The path to the file to check.
+
+    Returns:
+        A tuple containing a list of formatted error messages,
+        the extracted shapes dictionary, and the extracted scalars dictionary.
+    """
+    errors: list[str] = []
+    try:
+        code = filepath.read_text(encoding="utf-8")
+        lines = code.splitlines()
+        tree = ast.parse(code, filename=str(filepath))
+    except SyntaxError as e:
+        errors.append(f"{filepath}:{e.lineno}: [SyntaxError] {e.msg} ")
+        return errors, {}, {}
+
+    checker = Checker()
+    checker.visit(tree)
+
+    for warning in checker.warnings:
+        line = warning["line"]
+        line_idx = max(0, line - 1)
+        if line_idx < len(lines):
+            line_content = lines[line_idx]
+            if "# shpe: ignore" in line_content:
+                continue
+
+        code = warning["code"]
+        msg = warning["message"]
+        errors.append(f"{filepath}:{line}: warning: [{code}] {msg}")
+
+    for error in checker.errors:
+        line = error["line"]
+        line_idx = max(0, line - 1)
+        if line_idx < len(lines):
+            line_content = lines[line_idx]
+            if "# shpe: ignore" in line_content:
+                continue
+
+        code = error["code"]
+        msg = error["message"]
+        errors.append(f"{filepath}:{line}: [{code}] {msg}")
+
+    return errors, checker.env.shapes, checker.env.scalar_values
+
+
+def _print_shape_report(
+    all_shapes: dict[Path, dict[str, tuple[Any, ...] | ShapeState]],
+    all_scalars: dict[Path, dict[str, int | float | ShapeState]],
+) -> None:
+    """Prints a formatted report of all inferred shapes and scalar values across files.
+
+    Args:
+        all_shapes: Mapping of file paths to their symbol shape definitions.
+        all_scalars: Mapping of file paths to their tracked scalar variables.
+    """
+    print("\n-------- Symbol State (Shapes & Scalars) --------")
+
+    for filepath in sorted(all_shapes.keys(), key=str):
+        print(f"\n{filepath}:")
+        symbols = all_shapes[filepath]
+        scalars = all_scalars.get(filepath, {})
+
+        if scalars:
+            print("  Scalars:")
+            for name, val in sorted(scalars.items()):
+                val_str = "unknown" if val is ShapeState.UNKNOWN else val
+                print(f"    - {name} = {val_str}")
+
+        if symbols:
+            print("  Shapes:")
+            symbols_with_no_shape: list[str] = []
+            symbols_with_unknown_shape: list[str] = []
+            for var_name, shape in sorted(symbols.items()):
+                if shape is ShapeState.UNKNOWN:
+                    symbols_with_unknown_shape.append(var_name)
+                elif shape is not None:
+                    print(f"    - {var_name}: {shape}")
+                else:
+                    symbols_with_no_shape.append(var_name)
+
+            if symbols_with_unknown_shape:
+                print(
+                    "    - unknown shape (tainted) for: "
+                    f"{', '.join(symbols_with_unknown_shape)}"
+                )
+            if symbols_with_no_shape:
+                print(
+                    f"    - no shape inferred for: {', '.join(symbols_with_no_shape)}"
+                )
+
+        if not symbols and not scalars:
+            print("  - (no symbols tracked)")
+
+    print("\n" + "-" * 30 + "\n")
+
+
+if __name__ == "__main__":
+    main()

shpe/diagnostics.py

@@ -0,0 +1,85 @@
+import ast
+from enum import Enum
+from typing import Any
+
+
+class ErrorCode(Enum):
+    ANNOTATION = "Annotation"
+    ELEMENTWISE = "Elementwise"
+    MATMUL = "MatMul"
+    RESHAPE = "Reshape"
+    VALUE = "Value"
+
+
+def _position(node: ast.AST) -> tuple[int, int, int]:
+    col = getattr(node, "col_offset", 0)
+    return getattr(node, "lineno", 0), col, getattr(node, "end_col_offset", col)
+
+
+class Diagnostics:
+    """Collects errors, warnings, and inlay hints during a check pass."""
+
+    def __init__(self) -> None:
+        self.errors: list[dict[str, Any]] = []
+        self.warnings: list[dict[str, Any]] = []
+        self.inlay_hints: list[dict[str, Any]] = []
+        self.active_error_collection = True
+        self.active_warning_collection = True
+        self.active_hint_collection = True
+
+    def deactivate_hints(self) -> None:
+        """Deactivate inlay hints collection."""
+        self.active_hint_collection = False
+
+    def activate_hints(self) -> None:
+        """Activate inlay hints collection."""
+        self.active_hint_collection = True
+
+    def deactivate_errors(self) -> None:
+        """Deactivate error collection."""
+        self.active_error_collection = False
+
+    def activate_errors(self) -> None:
+        """Activate error collection."""
+        self.active_error_collection = True
+
+    def deactivate_warnings(self) -> None:
+        """Deactivate warning collection."""
+        self.active_warning_collection = False
+
+    def activate_warnings(self) -> None:
+        """Activate warning collection."""
+        self.active_warning_collection = True
+
+    def error(self, node: ast.AST, code: ErrorCode, message: str) -> None:
+        if not self.active_error_collection:
+            return
+        line, col, end_col = _position(node)
+        self.errors.append(
+            {
+                "line": line,
+                "col": col,
+                "end_col": end_col,
+                "code": code.value,
+                "message": message,
+            }
+        )
+
+    def warning(self, node: ast.AST, message: str) -> None:
+        if not self.active_warning_collection:
+            return
+        line, col, end_col = _position(node)
+        self.warnings.append(
+            {"line": line, "col": col, "end_col": end_col, "message": message}
+        )
+
+    def hint(self, node: ast.AST, shape: tuple[Any, ...]) -> None:
+        if not self.active_hint_collection:
+            return
+        self.inlay_hints.append(
+            {
+                "line": getattr(node, "lineno", 0),
+                "col": getattr(node, "end_col_offset", 0),
+                "shape": shape,
+            }
+        )

shpe/environment.py

@@ -0,0 +1,71 @@
+import ast
+from enum import Enum
+from typing import Any
+
+
+class ShapeState(Enum):
+    UNKNOWN = "UNKNOWN"
+
+
+class Environment:
+    def __init__(self) -> None:
+        self.stack: list[dict[str, Any]] = []
+        self.push_child()  # Initialize the root environment
+
+    @property
+    def shapes(self) -> dict[str, tuple[Any, ...] | ShapeState]:
+        return self.stack[-1]["shapes"]
+
+    @property
+    def scalar_values(self) -> dict[str, int | float | ShapeState]:
+        return self.stack[-1]["scalar_values"]
+
+    @property
+    def functions(self) -> dict[str, Any]:
+        return self.stack[-1]["functions"]
+
+    def __len__(self) -> int:
+        return len(self.stack)
+
+    def get_shape(self, name: str) -> tuple[Any, ...] | ShapeState | None:
+        for frame in reversed(self.stack):
+            if name in frame["shapes"]:
+                return frame["shapes"][name]
+        return None
+
+    def get_scalar(self, name: str) -> int | float | ShapeState | None:
+        for frame in reversed(self.stack):
+            if name in frame["scalar_values"]:
+                return frame["scalar_values"][name]
+        return None
+
+    def get_function(self, name: str) -> ast.FunctionDef | None:
+        for frame in reversed(self.stack):
+            if name in frame["functions"]:
+                return frame["functions"][name]
+        return None
+
+    def set_shape(self, name: str, shape: tuple[Any, ...] | ShapeState) -> None:
+        self.stack[-1]["shapes"][name] = shape
+
+    def set_scalar(self, name: str, value: int | float | ShapeState) -> None:
+        self.stack[-1]["scalar_values"][name] = value
+
+    def set_function(self, name: str, func: Any) -> None:
+        self.stack[-1]["functions"][name] = func
+
+    def push_child(self) -> None:
+        """Creates a new child environment for local variables and shapes."""
+        shapes: dict[str, tuple[Any, ...] | ShapeState] = {}
+        scalar_values: dict[str, int | float | ShapeState] = {}
+        functions: dict[str, ast.FunctionDef] = {}
+        self.stack.append(
+            {"shapes": shapes, "scalar_values": scalar_values, "functions": functions}
+        )
+
+    def pop_child(self) -> None:
+        """Removes the most recent child environment."""
+        if len(self.stack) > 1:
+            self.stack.pop()
+        else:
+            raise RuntimeError("Cannot pop the root environment.")