ending-locations: Fix bug for ClassDefs with no decorators

CHANGELOG.md

@@ -24,6 +24,7 @@ and adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   2. False negative when the loop variable can be simplified, but is also shadowed in the
      the loop body.
 - Fix HTML report to link correctly to specific errors on the PythonTA documentation website.
+- Fix bug when setting ending locations for `ClassDef`s that have no decorators.
 
 ### New checkers
 

examples/ending_locations/class_def.py

@@ -1,3 +1,8 @@
 @wrapper
 class Foo(base1, base2):
     pass
+
+
+class Bar:
+    def __init__(self):
+        pass

examples/ending_locations/function_def.py

@@ -5,3 +5,11 @@ def fun(arg) -> str:
     """
     return_annotation = "cool!"
     return return_annotation
+
+
+def fun2(arg) -> str:
+    """
+    This is a function fun2.
+    """
+    return_annotation = "cool!"
+    return return_annotation

python_ta/transforms/setendings.py

@@ -203,8 +203,8 @@ def init_register_ending_setters(source_code):
 
     # Ad hoc transformations
     ending_transformer.register_transform(nodes.BinOp, _set_start_from_first_child)
-    ending_transformer.register_transform(nodes.ClassDef, _set_start_from_first_child)
-    ending_transformer.register_transform(nodes.FunctionDef, _set_start_from_first_child)
+    ending_transformer.register_transform(nodes.ClassDef, _set_start_from_first_decorator)
+    ending_transformer.register_transform(nodes.FunctionDef, _set_start_from_first_decorator)
     ending_transformer.register_transform(nodes.Tuple, _set_start_from_first_child)
     ending_transformer.register_transform(nodes.Arguments, fix_arguments(source_code))
     ending_transformer.register_transform(nodes.Slice, fix_slice(source_code))
@@ -406,6 +406,15 @@ def _set_start_from_first_child(node):
     return node
 
 
+def _set_start_from_first_decorator(node):
+    """Set the start attributes of this node from its first child, if that child is a decorator."""
+    if getattr(node, "decorators"):
+        first_child = node.decorators
+        node.fromlineno = first_child.fromlineno
+        node.col_offset = first_child.col_offset
+    return node
+
+
 def set_from_last_child(node):
     """Populate ending locations for astroid node based on its last child.
 

tests/test_setendings.py

@@ -144,7 +144,7 @@ def test_call(self):
         self.set_and_check(module, nodes.Call, expected)
 
     def test_classdef(self):
-        expected = [(1, 3, 0, 8)]
+        expected = [(1, 3, 0, 8), (6, 8, 0, 12)]
         module = self.get_file_as_module("class_def.py")
         self.set_and_check(module, nodes.ClassDef, expected)
 
@@ -257,8 +257,8 @@ def test_for(self):
         self.set_and_check(module, nodes.For, expected)
 
     def test_functiondef(self):
-        """Note: this node includes the decorator."""
-        expected = [(1, 7, 0, 28)]
+        """Note: this node includes the decorator, if applicable."""
+        expected = [(1, 7, 0, 28), (10, 15, 0, 28)]
         module = self.get_file_as_module("function_def.py")
         self.set_and_check(module, nodes.FunctionDef, expected)