add class

Author: zhewenl

tools/stronghold/src/api/__init__.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import dataclasses
-from collections.abc import Sequence
+from collections.abc import Sequence, Mapping
 from typing import Optional
 
 import api.types
@@ -41,3 +41,30 @@ class Parameter:
     line: int
     # Type annotation (relies on ast.annotation types)
     type_annotation: Optional[api.types.TypeHint] = None
+
+
+@dataclasses.dataclass
+class Field:
+    """Represents a dataclass or class attribute."""
+
+    name: str
+    required: bool
+    line: int
+    type_annotation: Optional[api.types.TypeHint] = None
+
+
+@dataclasses.dataclass
+class Class:
+    """Represents a class or dataclass."""
+
+    fields: Sequence[Field]
+    line: int
+    dataclass: bool = False
+
+
+@dataclasses.dataclass
+class API:
+    """Represents extracted API information."""
+
+    functions: Mapping[str, Parameters]
+    classes: Mapping[str, Class]

tools/stronghold/src/api/ast.py

@@ -11,28 +11,27 @@
 import api.types
 
 
-def extract(path: pathlib.Path) -> Mapping[str, api.Parameters]:
-    """Extracts the API from a given source file.
-
-    The keys will be the fully-qualified path from the root of the module, e.g.
-     * global_func
-     * ClassName.method_name
-     * ClassName.SubClassName.method_name
-    """
-    raw_api = extract_raw(path)
-    return {
-        name: _function_def_to_parameters(function_def)
-        for name, function_def in raw_api.items()
+def extract(path: pathlib.Path, *, include_classes: bool = False) -> api.API:
+    """Extracts API definitions from a given source file."""
+
+    funcs, classes = extract_raw(path, include_classes=include_classes)
+    parameters = {
+        name: _function_def_to_parameters(func) for name, func in funcs.items()
     }
+    return api.API(functions=parameters, classes=classes)
+
 
+def extract_raw(
+    path: pathlib.Path, *, include_classes: bool = False
+) -> tuple[Mapping[str, ast.FunctionDef], Mapping[str, api.Class]]:
+    """Extracts API as AST nodes."""
 
-def extract_raw(path: pathlib.Path) -> Mapping[str, ast.FunctionDef]:
-    """Extracts the API as ast.FunctionDef instances."""
-    out: dict[str, ast.FunctionDef] = {}
-    _ContextualNodeVisitor(out, context=[]).visit(
+    funcs: dict[str, ast.FunctionDef] = {}
+    classes: dict[str, api.Class] = {}
+    _ContextualNodeVisitor(funcs, classes if include_classes else None, []).visit(
         ast.parse(path.read_text(), os.fspath(path))
     )
-    return out
+    return funcs, classes
 
 
 def _function_def_to_parameters(node: ast.FunctionDef) -> api.Parameters:
@@ -90,20 +89,69 @@ def _function_def_to_parameters(node: ast.FunctionDef) -> api.Parameters:
 
 
 class _ContextualNodeVisitor(ast.NodeVisitor):
-    """NodeVisitor implementation that tracks which class, if any, it is a member of."""
-
-    def __init__(self, out: dict[str, ast.FunctionDef], context: Sequence[str]) -> None:
-        self._out = out
-        self._context = context
+    """NodeVisitor that collects functions and optionally classes."""
+
+    def __init__(
+        self,
+        functions: dict[str, ast.FunctionDef],
+        classes: dict[str, api.Class] | None,
+        context: Sequence[str],
+    ) -> None:
+        self._functions = functions
+        self._classes = classes
+        self._context = list(context)
 
     def visit_ClassDef(self, node: ast.ClassDef) -> None:
         # Recursively visit all nodes under this class, with the given
         # class name pushed onto a new context.
+        if self._classes is not None:
+            name = ".".join(self._context + [node.name])
+            is_dataclass = any(
+                (isinstance(dec, ast.Name) and dec.id == "dataclass")
+                or (isinstance(dec, ast.Attribute) and dec.attr == "dataclass")
+                for dec in node.decorator_list
+            )
+            fields: list[api.Field] = []
+            for stmt in node.body:
+                if isinstance(stmt, ast.AnnAssign) and isinstance(
+                    stmt.target, ast.Name
+                ):
+                    field_name = stmt.target.id
+                    if field_name.startswith("_"):
+                        continue
+                    fields.append(
+                        api.Field(
+                            name=field_name,
+                            required=stmt.value is None,
+                            line=stmt.lineno,
+                            type_annotation=api.types.annotation_to_dataclass(
+                                stmt.annotation
+                            ),
+                        )
+                    )
+                elif isinstance(stmt, ast.Assign):
+                    for target in stmt.targets:
+                        if isinstance(target, ast.Name):
+                            field_name = target.id
+                            if field_name.startswith("_"):
+                                continue
+                            fields.append(
+                                api.Field(
+                                    name=field_name,
+                                    required=False,
+                                    line=stmt.lineno,
+                                    type_annotation=None,
+                                )
+                            )
+            self._classes[name] = api.Class(
+                fields=fields, line=node.lineno, dataclass=is_dataclass
+            )
+
         _ContextualNodeVisitor(
-            self._out, list(self._context) + [node.name]
+            self._functions, self._classes, self._context + [node.name]
         ).generic_visit(node)
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
         # Records this function.
-        name = ".".join(list(self._context) + [node.name])
-        self._out[name] = node
+        name = ".".join(self._context + [node.name])
+        self._functions[name] = node

tools/stronghold/src/api/compatibility.py

@@ -68,15 +68,19 @@ def check(
     before: pathlib.Path, after: pathlib.Path
 ) -> Sequence[api.violations.Violation]:
     """Identifies API compatibility issues between two files."""
-    before_api = api.ast.extract(before)
-    after_api = api.ast.extract(after)
+    before_api = api.ast.extract(before, include_classes=True)
+    after_api = api.ast.extract(after, include_classes=True)
+    before_funcs = before_api.functions
+    after_funcs = after_api.functions
+    before_classes = before_api.classes
+    after_classes = after_api.classes
 
     violations: list[api.violations.Violation] = []
-    for name, before_def in before_api.items():
+    for name, before_def in before_funcs.items():
         if any(token.startswith("_") for token in name.split(".")):
             continue
 
-        after_def = after_api.get(name)
+        after_def = after_funcs.get(name)
         if after_def is None:
             violations.append(api.violations.FunctionDeleted(func=name, line=1))
             continue
@@ -103,6 +107,14 @@ def check(
         violations += _check_by_requiredness(name, before_def, after_def)
         violations += _check_variadic_parameters(name, before_def, after_def)
 
+    for name, before_class in before_classes.items():
+        if any(token.startswith("_") for token in name.split(".")):
+            continue
+        after_class = after_classes.get(name)
+        if after_class is None:
+            continue
+        violations += list(_check_class_fields(name, before_class, after_class))
+
     return violations
 
 
@@ -250,6 +262,37 @@ def _check_variadic_parameters(
         yield api.violations.KwArgsDeleted(func, line=after.line)
 
 
+def _check_class_fields(
+    cls: str, before: api.Class, after: api.Class
+) -> Iterable[api.violations.Violation]:
+    """Checks class and dataclass field compatibility."""
+
+    before_fields = {f.name: f for f in before.fields}
+    after_fields = {f.name: f for f in after.fields}
+
+    for name, before_field in before_fields.items():
+        after_field = after_fields.get(name)
+        if after_field is None:
+            yield api.violations.FieldRemoved(func=cls, parameter=name, line=after.line)
+            continue
+            
+        if not _check_type_compatibility(
+            before_field.type_annotation, after_field.type_annotation
+        ):
+            yield api.violations.FieldTypeChanged(
+                func=cls,
+                parameter=name,
+                line=after_field.line,
+                type_before=str(before_field.type_annotation),
+                type_after=str(after_field.type_annotation),
+            )
+
+    for name in set(after_fields) - set(before_fields):
+        yield api.violations.FieldAdded(
+            func=cls, parameter=name, line=after_fields[name].line
+        )
+
+
 def _check_type_compatibility(
     type_before: api.types.TypeHint, type_after: api.types.TypeHint
 ) -> bool:

tools/stronghold/src/api/violations.py

@@ -123,3 +123,38 @@ def __post_init__(self) -> None:
         self.message = (
             f"{self.parameter} changed from {self.type_before} to {self.type_after}"
         )
+
+
+# ====================================
+# Class field violations
+@dataclass
+class FieldViolation(Violation):
+    parameter: str = ""
+
+
+@dataclass
+class FieldRemoved(FieldViolation):
+    message: str = ""
+
+    def __post_init__(self) -> None:
+        self.message = f"{self.parameter} was removed"
+
+
+@dataclass
+class FieldAdded(FieldViolation):
+    message: str = ""
+
+    def __post_init__(self) -> None:
+        self.message = f"{self.parameter} was added"
+
+
+@dataclass
+class FieldTypeChanged(FieldViolation):
+    type_before: str = ""
+    type_after: str = ""
+    message: str = ""
+
+    def __post_init__(self) -> None:
+        self.message = (
+            f"{self.parameter} changed from {self.type_before} to {self.type_after}"
+        )