Merge pull request #83 from TimRepke/main

Fix handling of decorated classes

Author: martinjm97

tap/utils.py

@@ -190,8 +190,11 @@ def tokenize_source(obj: object) -> Generator:
 
 def get_class_column(obj: type) -> int:
     """Determines the column number for class variables in a class."""
+    first_line = 1
     for token_type, token, (start_line, start_column), (end_line, end_column), line in tokenize_source(obj):
-        if start_line == 1 or token.strip() == '':
+        if token.strip() == '@':
+            first_line += 1
+        if start_line <= first_line or token.strip() == '':
             continue
 
         return start_column

tests/test_utils.py

@@ -137,6 +137,10 @@ def test_type_to_str(self) -> None:
         self.assertEqual(type_to_str(Union[List[int], Dict[float, bool]]), 'Union[List[int], Dict[float, bool]]')
 
 
+def class_decorator(cls):
+    return cls
+
+
 class ClassColumnTests(TestCase):
     def test_column_simple(self):
         class SimpleColumn:
@@ -167,6 +171,23 @@ def func(self):
 
         self.assertEqual(get_class_column(FuncColumn), 12)
 
+    def test_dataclass(self):
+        @class_decorator
+        class DataclassColumn:
+            arg: int = 5
+        self.assertEqual(get_class_column(DataclassColumn), 12)
+
+    def test_dataclass_method(self):
+        def wrapper(f):
+            pass
+
+        @class_decorator
+        class DataclassColumn:
+            @wrapper
+            def func(self):
+                pass
+        self.assertEqual(get_class_column(DataclassColumn), 12)
+
 
 class ClassVariableTests(TestCase):
     def test_no_variables(self):
@@ -282,6 +303,14 @@ def f(self):
         class_variables['i'] = {'comment': ''}
         self.assertEqual(get_class_variables(FunctionsWithDocs), class_variables)
 
+    def test_dataclass(self):
+        @class_decorator
+        class DataclassColumn:
+            arg: int = 5
+        class_variables = OrderedDict()
+        class_variables['arg'] = {'comment': ''}
+        self.assertEqual(get_class_variables(DataclassColumn), class_variables)
+
 
 class GetLiteralsTests(TestCase):
     def test_get_literals_string(self) -> None: