Class statement support (#22)

fix: #18

setup.py

@@ -16,10 +16,11 @@
 # under the License.
 
 """The script for setting up stmdency."""
+from __future__ import annotations
+
 import logging
 import os
 from distutils.dir_util import remove_tree
-from typing import List
 
 from setuptools import Command, setup
 
@@ -30,7 +31,7 @@ class CleanCommand(Command):
  """Command to clean up python api before setup by running `python setup.py pre_clean`."""
 
  description = "Clean up project root"
- user_options: List[str] = []
+ user_options: list[str] = []
  clean_list = [
  "build",
  "htmlcov",

src/stmdency/visitors/base.py

@@ -4,7 +4,7 @@
 
 import libcst as cst
 import libcst.matchers as m
-from libcst import Assign, FunctionDef, Import, ImportFrom
+from libcst import Assign, ClassDef, FunctionDef, Import, ImportFrom
 
 from stmdency.models.node import StmdencyNode
 from stmdency.visitors.assign import AssignVisitor
@@ -22,7 +22,7 @@ class BaseVisitor(cst.CSTVisitor):
 
  stack: dict[str, StmdencyNode] = field(default_factory=dict)
  # Add scope to determine if the node is in the same scope
- scope: dict[cst.CSTNode] = field(default_factory=set)
+ scope: set[cst.CSTNode] = field(default_factory=set)
 
  def handle_import(self, node: Import | ImportFrom) -> None:
  """Handle `import` / `from xx import xxx` statement and parse/add to stack."""
@@ -44,6 +44,19 @@ def visit_ImportFrom(self, node: ImportFrom) -> bool | None:
  self.handle_import(node)
  return True
 
+ def visit_ClassDef(self, node: ClassDef) -> bool | None:
+ """Handle class definition, pass to ClassDefVisitor and add scope.
+
+ the reason add scope is to skip the visit_Assign in current class
+ """
+ self.scope.add(node)
+ self.stack.update([(node.name.value, StmdencyNode(node=node))])
+ return True
+
+ def leave_ClassDef(self, original_node: ClassDef) -> None:
+ """Remove class definition in scope."""
+ self.scope.remove(original_node)
+
  def visit_FunctionDef(self, node: FunctionDef) -> bool | None:
  """Handle function definition, pass to FunctionVisitor and add scope.
 

tests/extractor/test_assign.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from __future__ import annotations
 
 import pytest
 
@@ -81,5 +81,5 @@
 
 
 @pytest.mark.parametrize("name, source, expects", assign_cases)
-def test_assign(name: str, source: str, expects: Dict[str, str]) -> None:
+def test_assign(name: str, source: str, expects: dict[str, str]) -> None:
  assert_extract(name, source, expects)

tests/extractor/test_function.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from __future__ import annotations
 
 import pytest
 
@@ -296,5 +296,5 @@ def foo():
 
 
 @pytest.mark.parametrize("name, source, expects", func_cases)
-def test_func(name: str, source: str, expects: Dict[str, str]) -> None:
+def test_func(name: str, source: str, expects: dict[str, str]) -> None:
  assert_extract(name, source, expects)

tests/extractor/test_import.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from __future__ import annotations
 
 import pytest
 
@@ -124,5 +124,5 @@ def bar():
 
 
 @pytest.mark.parametrize("name, source, expects", import_cases)
-def test_import(name: str, source: str, expects: Dict[str, str]) -> None:
+def test_import(name: str, source: str, expects: dict[str, str]) -> None:
  assert_extract(name, source, expects)

tests/extractor/test_module_class.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+import pytest
+
+from tests.testing import assert_extract
+
+import_cases = [
+ (
+ "simple class with init",
+ """
+ class Foo:
+ def __init__(self, arg1):
+ self.arg1 = arg1
+ def foo():
+ f = Foo(arg1=1)
+ print(f)
+ """,
+ {
+ "foo": """\
+ class Foo:
+ def __init__(self, arg1):
+ self.arg1 = arg1\n\n
+ def foo():
+ f = Foo(arg1=1)
+ print(f)
+ """,
+ },
+ ),
+]
+
+
+@pytest.mark.parametrize("name, source, expects", import_cases)
+def test_import(name: str, source: str, expects: dict[str, str]) -> None:
+ assert_extract(name, source, expects)

tests/testing.py

@@ -1,10 +1,11 @@
+from __future__ import annotations
+
 import textwrap
-from typing import Dict
 
 from stmdency.extractor import Extractor
 
 
-def assert_extract(name: str, source: str, expects: Dict[str, str]) -> None:
+def assert_extract(name: str, source: str, expects: dict[str, str]) -> None:
  wrap_source = textwrap.dedent(source)
  extractor = Extractor(source=wrap_source)
  for expect in expects: