Add support for class attributes

semver_bumper/git.py

@@ -12,7 +12,6 @@
     "find_file_contents_in_ref",
     "find_initial_commit_hash",
     "find_previous_version",
-    "get_head_commit_hash",
     "has_new_commits_since_ref",
     "run_command",
 ]

semver_bumper/parser.py

@@ -2,9 +2,19 @@
 
 import ast
 import json
-
-from .typing import ArgKind, ArgumentData, AssignmentData, BodyData, ClassData, FunctionArguments, FunctionData
-from .utils import is_dunder_method, is_internal_method
+from typing import Iterable
+
+from .typing import (
+    ArgKind,
+    ArgumentData,
+    AssignmentData,
+    BodyData,
+    ClassAttributeData,
+    ClassData,
+    FunctionArguments,
+    FunctionData,
+)
+from .utils import is_class_dunder_method, is_class_internal_method, is_dunder_method, is_internal_method
 
 __all__ = [
     "parse_module_body",
@@ -49,6 +59,7 @@ def _parse_body(nodes: list[ast.stmt], *, dunder_all: list[str] | None = None) -
             body.classes[name] = ClassData(
                 name=name,
                 body=_parse_body(node.body),
+                attributes=_get_class_attributes(node.body),
             )
 
         elif isinstance(node, ast.Assign) and _should_include_assignment(node, dunder_all):
@@ -112,23 +123,44 @@ def _get_argument_data(node: ast.FunctionDef) -> list[ArgumentData]:
     return data
 
 
-def _should_include_function(node: ast.FunctionDef, dunder_all: list[str] | None) -> bool:
+def _get_class_attributes(nodes: list[ast.stmt]) -> dict[str, ClassAttributeData]:
+    """Get class attributes from the body of a class definition."""
+    data: dict[str, ClassAttributeData] = {}
+    for node in nodes:
+        # Only look for class attribute declarations from the `__init__` method.
+        if not isinstance(node, ast.FunctionDef) or node.name != "__init__":
+            continue
+
+        args = _get_argument_data(node)
+        for item in node.body:
+            if isinstance(item, ast.Assign) and _should_include_class_assignment(item):
+                name = _node_to_string(item.targets[0]).split(".")[-1]
+                data[name] = ClassAttributeData(name=name, type=_infer_type_for_node(item.value, args))
+
+            elif isinstance(item, ast.AnnAssign) and _should_include_class_ann_assignment(item):
+                name = _node_to_string(item.target).split(".")[-1]
+                data[name] = ClassAttributeData(name=name, type=_node_to_string(item.annotation))
+
+    return data
+
+
+def _should_include_function(node: ast.FunctionDef, dunder_all: list[str] | None = None) -> bool:
     """Should the given function definition be included in the diffing data?"""
     name = node.name
     if dunder_all is not None:
         return name in dunder_all
     return not is_internal_method(name) or is_dunder_method(name)
 
 
-def _should_include_class(node: ast.ClassDef, dunder_all: list[str] | None) -> bool:
+def _should_include_class(node: ast.ClassDef, dunder_all: list[str] | None = None) -> bool:
     """Should the given class definition be included in the diffing data?"""
     name = node.name
     if dunder_all is not None:
         return name in dunder_all
     return not is_internal_method(name) or is_dunder_method(name)
 
 
-def _should_include_assignment(node: ast.Assign, dunder_all: list[str] | None) -> bool:
+def _should_include_assignment(node: ast.Assign, dunder_all: list[str] | None = None) -> bool:
     """Should the given assignment be included in the diffing data?"""
     name = _node_to_string(node.targets[0])
     if name == "__all__":
@@ -138,29 +170,47 @@ def _should_include_assignment(node: ast.Assign, dunder_all: list[str] | None) -
     return not is_internal_method(name) or is_dunder_method(name)
 
 
-def _should_include_ann_assignment(node: ast.AnnAssign, dunder_all: list[str] | None) -> bool:
+def _should_include_ann_assignment(node: ast.AnnAssign, dunder_all: list[str] | None = None) -> bool:
     """Should the given annotated assignment be included in the diffing data?"""
     name = _node_to_string(node.target)
     if dunder_all is not None:
         return name in dunder_all
     return not is_internal_method(name) or is_dunder_method(name)
 
 
+def _should_include_class_assignment(node: ast.Assign) -> bool:
+    """Should the given class assignment be included in the diffing data?"""
+    name = _node_to_string(node.targets[0])
+    return not is_class_internal_method(name) or is_class_dunder_method(name)
+
+
+def _should_include_class_ann_assignment(node: ast.AnnAssign) -> bool:
+    """Should the given class assignment be included in the diffing data?"""
+    name = _node_to_string(node.target)
+    return not is_class_internal_method(name) or is_class_dunder_method(name)
+
+
 def _node_to_string(node: ast.expr | None) -> str | None:
     """Convert the given ast node to a string."""
     if node is None:
         return None
     return ast.unparse(node)
 
 
-def _infer_type_for_node(node: ast.expr | None) -> str | None:  # noqa: PLR0911
+def _infer_type_for_node(node: ast.expr | None, args: Iterable[ArgumentData] = ()) -> str | None:  # noqa: PLR0911
     """Infer the type for the given ast node."""
     if node is None:
         return None
 
     if isinstance(node, ast.Constant):
         return type(node.value).__name__
 
+    if isinstance(node, ast.Name):
+        for arg in args:
+            if arg.name == node.id:
+                return arg.type
+        return None
+
     if isinstance(node, ast.List):
         subtypes = _infer_types_for_nodes(node.elts)
         return f"list[{subtypes}]"

semver_bumper/typing.py

@@ -37,6 +37,12 @@ class ArgumentData:
     kind: ArgKind
 
 
+@dataclasses.dataclass
+class ClassAttributeData:
+    name: str
+    type: str | None
+
+
 @dataclasses.dataclass
 class FunctionData:
     name: str
@@ -48,6 +54,7 @@ class FunctionData:
 class ClassData:
     name: str
     body: BodyData
+    attributes: dict[str, ClassAttributeData] = dataclasses.field(default_factory=dict)
 
 
 @dataclasses.dataclass

semver_bumper/utils.py

@@ -47,6 +47,22 @@ def is_dunder_method(name: str) -> bool:
     return name.startswith("__") and name.endswith("__")
 
 
+def is_class_internal_method(name: str) -> bool:
+    """Is the given name an internal method?"""
+    parts = name.split(".")
+    if len(parts) != 2:  # noqa: PLR2004
+        return False
+    return parts[0] == "self" and is_internal_method(parts[1])
+
+
+def is_class_dunder_method(name: str) -> bool:
+    """Is the given name an internal method?"""
+    parts = name.split(".")
+    if len(parts) != 2:  # noqa: PLR2004
+        return False
+    return parts[0] == "self" and is_dunder_method(parts[1])
+
+
 def find_python_files(path: Path) -> Generator[Path, None, None]:
     """
     Find all python files in the given directory and its subdirectories.

tests/example/arg_spec_functions.py

@@ -480,7 +480,8 @@ class Class_1:
     class Class_2:
         foo = {1: "foo", "z": [1, 2, 3]}
 
-    def __init__(self, foo: Foo) -> None: ...
+    def __init__(self, bar: Bar) -> None:
+        self.bar = bar
 
     def method(self, foo: int) -> None: ...
 

tests/test_parse_file.py

@@ -38,6 +38,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Bar"] == {
         "name": "Bar",
@@ -59,6 +60,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Barr"] == {
         "name": "Barr",
@@ -80,6 +82,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Baz"] == {
         "name": "Baz",
@@ -101,6 +104,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Bazz"] == {
         "name": "Bazz",
@@ -122,6 +126,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Fizz"] == {
         "name": "Fizz",
@@ -139,6 +144,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Buzz"] == {
         "name": "Buzz",
@@ -152,6 +158,7 @@ def test_parse_file() -> None:
             "classes": {},
             "functions": {},
         },
+        "attributes": {},
     }
     assert results["classes"]["Class_1"] == {
         "name": "Class_1",
@@ -175,6 +182,7 @@ def test_parse_file() -> None:
                         "classes": {},
                         "functions": {},
                     },
+                    "attributes": {},
                 }
             },
             "functions": {
@@ -187,8 +195,8 @@ def test_parse_file() -> None:
                             "kind": ArgKind.REGULAR,
                         },
                         {
-                            "name": "foo",
-                            "type": "Foo",
+                            "name": "bar",
+                            "type": "Bar",
                             "kind": ArgKind.REGULAR,
                         },
                     ],
@@ -250,6 +258,12 @@ def test_parse_file() -> None:
                 },
             },
         },
+        "attributes": {
+            "bar": {
+                "name": "bar",
+                "type": "Bar",
+            },
+        },
     }
 
     name = func_name()