Fix handling of decorated classes

This should resolve the issue described in #80
where '@' decorated classes are missing attributes/arguments due to misaligned
class column during tokenisation.

Author: TimRepke

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

@@ -1,5 +1,6 @@
 from argparse import ArgumentTypeError
 from collections import OrderedDict
+from dataclasses import dataclass
 import json
 import os
 import subprocess
@@ -167,6 +168,23 @@ def func(self):
 
         self.assertEqual(get_class_column(FuncColumn), 12)
 
+    def test_dataclass(self):
+        @dataclass
+        class DataclassColumn:
+            arg: int = 5
+        self.assertEqual(get_class_column(DataclassColumn), 12)
+
+    def test_dataclass_method(self):
+        def wrapper(f):
+            pass
+
+        @dataclass
+        class DataclassColumn:
+            @wrapper
+            def func(self):
+                pass
+        self.assertEqual(get_class_column(DataclassColumn), 12)
+
 
 class ClassVariableTests(TestCase):
     def test_no_variables(self):
@@ -282,6 +300,14 @@ def f(self):
         class_variables['i'] = {'comment': ''}
         self.assertEqual(get_class_variables(FunctionsWithDocs), class_variables)
 
+    def test_dataclass(self):
+        @dataclass
+        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: