[RFC] add class support in bc-linter (#6953)

as discussed with @izaitsevfb, we want to also include linter checks on
classes in addition to public functions; but prior to landing this PR,
we need to create some rules template for linters, where we can define
what to check(or not).
Like:
```
pytorch:
- Files: global
- Include: functions

vLLM:
- Files: vllm/v1/attrition, vllm/v1/core
- Include, classes, functions
```


cc @huydhn @yangw-dev @houseroad

---------

Co-authored-by: Ivan Zaitsev <[email protected]>

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 Mapping, Sequence
 from typing import Optional
 
 import api.types
@@ -41,3 +41,32 @@ 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
+    # Whether the field participates in the dataclass __init__ (dataclasses only)
+    init: bool = True
+
+
+@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,105 @@ 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
+                    required = stmt.value is None
+                    init = True
+                    # Support dataclasses.field(...)
+                    if isinstance(stmt.value, ast.Call):
+                        fn = stmt.value.func
+
+                        def _is_field_func(f: ast.AST) -> bool:
+                            return (isinstance(f, ast.Name) and f.id == "field") or (
+                                isinstance(f, ast.Attribute) and f.attr == "field"
+                            )
+
+                        if _is_field_func(fn):
+                            # default/default_factory imply not required
+                            has_default = any(
+                                isinstance(kw, ast.keyword)
+                                and kw.arg == "default"
+                                and kw.value is not None
+                                for kw in stmt.value.keywords
+                            )
+                            has_default_factory = any(
+                                isinstance(kw, ast.keyword)
+                                and kw.arg == "default_factory"
+                                and kw.value is not None
+                                for kw in stmt.value.keywords
+                            )
+                            required = not (has_default or has_default_factory)
+                            # init flag
+                            for kw in stmt.value.keywords:
+                                if isinstance(kw, ast.keyword) and kw.arg == "init":
+                                    init = not (
+                                        isinstance(kw.value, ast.Constant)
+                                        and kw.value.value is False
+                                    )
+                                    break
+                    fields.append(
+                        api.Field(
+                            name=field_name,
+                            required=required,
+                            line=stmt.lineno,
+                            type_annotation=api.types.annotation_to_dataclass(
+                                stmt.annotation
+                            ),
+                            init=init,
+                        )
+                    )
+                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,
+                                    init=True,
+                                )
+                            )
+            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,39 @@ 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
+
+    # Identify deleted classes to avoid double-reporting their methods as deleted
+    deleted_classes = {
+        name
+        for name in before_classes
+        if not any(token.startswith("_") for token in name.split("."))
+        and name not in after_classes
+    }
+
+    def _under_deleted_class(func_name: str) -> bool:
+        tokens = func_name.split(".")
+        prefix = []
+        for t in tokens[:-1]:  # exclude the function name itself
+            prefix.append(t)
+            if ".".join(prefix) in deleted_classes:
+                return True
+        return False
 
     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
+        if _under_deleted_class(name):
+            # Will be reported as a class deletion instead
+            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 +127,17 @@ 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))
+
+    # Classes deleted between before and after
+    violations += list(_check_deleted_classes(before_classes, after_classes))
+
     return violations
 
 
@@ -247,7 +282,68 @@ def _check_variadic_parameters(
     if before.variadic_args and not after.variadic_args:
         yield api.violations.VarArgsDeleted(func=func, line=after.line)
     if before.variadic_kwargs and not after.variadic_kwargs:
-        yield api.violations.KwArgsDeleted(func, line=after.line)
+        yield api.violations.KwArgsDeleted(func=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):
+        # Policy: Only flag additions for dataclasses, and only when required and in __init__
+        if after.dataclass:
+            added = after_fields[name]
+            if added.required and getattr(added, "init", True):
+                yield api.violations.FieldAdded(
+                    func=cls, parameter=name, line=added.line
+                )
+
+
+def _check_deleted_classes(
+    before_classes: Mapping[str, api.Class], after_classes: Mapping[str, api.Class]
+) -> Iterable[api.violations.Violation]:
+    """Emits violations for classes deleted between before and after."""
+    deleted = [
+        name
+        for name in before_classes
+        if not any(token.startswith("_") for token in name.split("."))
+        and name not in after_classes
+    ]
+
+    deleted_set = set(deleted)
+
+    def has_deleted_ancestor(class_name: str) -> bool:
+        tokens = class_name.split(".")
+        for i in range(1, len(tokens)):
+            if ".".join(tokens[:i]) in deleted_set:
+                return True
+        return False
+
+    for name in deleted:
+        if not has_deleted_ancestor(name):
+            # Align with FunctionDeleted's use of line=1
+            yield api.violations.ClassDeleted(func=name, line=1)
 
 
 def _check_type_compatibility(

tools/stronghold/src/api/violations.py

@@ -123,3 +123,47 @@ 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}"
+        )
+
+
+# ====================================
+# Class violations
+@dataclass
+class ClassDeleted(Violation):
+    """Represents a public class being deleted."""
+
+    message: str = "class deleted"