From 51aa20955ef674a24ea69ef7971e407e608853d4 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 16:13:41 +0530
Subject: [PATCH 01/17] Use ast module to find class definition

---
 Lib/inspect.py              | 71 +++++++++++++++++++++++++++----------
 Lib/test/inspect_fodder2.py | 71 ++++++++++++++++++++++++++++++++++---
 Lib/test/test_inspect.py    | 16 +++++++++
 3 files changed, 135 insertions(+), 23 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 6f7d5cd19ce4c6..4ba44598f96b82 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -802,25 +802,58 @@ def findsource(object):
         return lines, 0
 
     if isclass(object):
-        name = object.__name__
-        pat = re.compile(r'^(\s*)class\s*' + name + r'\b')
-        # make some effort to find the best matching class definition:
-        # use the one with the least indentation, which is the one
-        # that's most probably not inside a function definition.
-        candidates = []
-        for i in range(len(lines)):
-            match = pat.match(lines[i])
-            if match:
-                # if it's at toplevel, it's already the best one
-                if lines[i][0] == 'c':
-                    return lines, i
-                # else add whitespace to candidate list
-                candidates.append((match.group(1), i))
-        if candidates:
-            # this will sort by whitespace, and by line number,
-            # less whitespace first
-            candidates.sort()
-            return lines, candidates[0][1]
+        import ast
+
+        class ClassVisitor(ast.NodeVisitor):
+
+            def recursive(func):
+                def wrapper(self, node):
+                    """A recursive wrapper to iterate over inner classes that uses a stack
+                    to store inner class's name for computation of qualname"""
+                    if not self.stack:
+                        self.stack.append(node.name)
+
+                    func(self, node)
+                    for child in ast.iter_child_nodes(node):
+                        if isinstance(child, ast.ClassDef):
+                            self.stack.append(child.name)
+                            self.visit(child)
+                            self.stack.pop()
+
+                    if self.stack:
+                        self.stack.pop()
+
+                return wrapper
+
+            def __init__(self, source, qualname, *args, **kwargs):
+                self.stack = []
+                self.source = source
+                self.line_number = None
+                self.qualname = qualname
+                super().__init__(*args, **kwargs)
+
+            @recursive
+            def visit_ClassDef(self, node):
+                qualname = '.'.join(self.stack)
+                if self.qualname == qualname:
+                    # decrement by one since lines starts with indexing by zero
+                    self.line_number = node.lineno - 1
+
+        qualname = object.__qualname__
+        source = ''.join(lines)
+        tree = ast.parse(source)
+        class_visitor = ClassVisitor(source, qualname)
+        class_visitor.visit(tree)
+
+        if class_visitor.line_number is not None:
+            line_number = class_visitor.line_number
+            decorator_pattern = re.compile(r'^(\s*@)')
+            for line in reversed(lines[:line_number]):
+                if decorator_pattern.match(line):
+                    line_number -= 1
+                else:
+                    break
+            return lines, line_number
         else:
             raise OSError('could not find class definition')
 
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index 5a7b559d07d762..24777a57fd21b2 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -138,18 +138,81 @@ def func137():
             never_reached1
             never_reached2
 
-#line 141
+# line 141
+class cls142:
+    a = """
+class cls149:
+    ...
+"""
+
+# line 148
+class cls149:
+
+    def func151(self):
+        pass
+
+'''
+class cls160:
+    pass
+'''
+
+# line 159
+class cls160:
+
+    def func162(self):
+        pass
+
+# line 165
+class cls166:
+    a = '''
+    class cls175:
+        ...
+    '''
+
+# line 172
+class cls173:
+
+    class cls175:
+        pass
+
+# line 178
+class cls179:
+    pass
+
+# line 182
+class cls183:
+
+    class cls185:
+
+        def func186(self):
+            pass
+
+
+def class_decorator(cls):
+    return cls
+
+# line 195
+@class_decorator
+@class_decorator
+class cls197:
+
+    @class_decorator
+    @class_decorator
+    class cls201:
+        pass
+
+#line 204
 def positional_only_arg(a, /):
     pass
 
-#line 145
+#line 208
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 149
+# line 212
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 153
+#line 216
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 2dc8454595e175..82b256492c2f33 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -709,6 +709,22 @@ def test_getsource_on_method(self):
     def test_nested_func(self):
         self.assertSourceEqual(mod2.cls135.func136, 136, 139)
 
+    def test_class_definition_in_multiline_string_definition(self):
+        self.assertSourceEqual(mod2.cls149, 149, 152)
+
+    def test_class_definition_in_multiline_comment(self):
+        self.assertSourceEqual(mod2.cls160, 160, 163)
+
+    def test_nested_class_definition_indented_string(self):
+        self.assertSourceEqual(mod2.cls173.cls175, 175, 176)
+
+    def test_nested_class_definition(self):
+        self.assertSourceEqual(mod2.cls183, 183, 188)
+        self.assertSourceEqual(mod2.cls183.cls185, 185, 188)
+
+    def test_class_decorator(self):
+        self.assertSourceEqual(mod2.cls197, 195, 202)
+        self.assertSourceEqual(mod2.cls197.cls201, 199, 202)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From 341933df107bb093b359194f1cb544c686ef0936 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 16:23:44 +0530
Subject: [PATCH 02/17] Add NEWS entry

---
 .../next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst      | 3 +++
 1 file changed, 3 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst

diff --git a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
new file mode 100644
index 00000000000000..8919ad663493b9
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
@@ -0,0 +1,3 @@
+:meth:`inspect.getsource` now returns correct source code for inner class
+with same name as module level class. class decorators are also returned as
+part of the source.

From dbbb38ee98d97aa6d93a6999c54d37a8499824ac Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 17:45:05 +0530
Subject: [PATCH 03/17] Fix class with multiple children and move decorator
 code to the method

---
 Lib/inspect.py              | 39 ++++++++++++++++++-------------------
 Lib/test/inspect_fodder2.py | 16 +++++++++++----
 Lib/test/test_inspect.py    |  7 +++++++
 3 files changed, 38 insertions(+), 24 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 4ba44598f96b82..b1705a5284d3ee 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -806,25 +806,6 @@ def findsource(object):
 
         class ClassVisitor(ast.NodeVisitor):
 
-            def recursive(func):
-                def wrapper(self, node):
-                    """A recursive wrapper to iterate over inner classes that uses a stack
-                    to store inner class's name for computation of qualname"""
-                    if not self.stack:
-                        self.stack.append(node.name)
-
-                    func(self, node)
-                    for child in ast.iter_child_nodes(node):
-                        if isinstance(child, ast.ClassDef):
-                            self.stack.append(child.name)
-                            self.visit(child)
-                            self.stack.pop()
-
-                    if self.stack:
-                        self.stack.pop()
-
-                return wrapper
-
             def __init__(self, source, qualname, *args, **kwargs):
                 self.stack = []
                 self.source = source
@@ -832,13 +813,31 @@ def __init__(self, source, qualname, *args, **kwargs):
                 self.qualname = qualname
                 super().__init__(*args, **kwargs)
 
-            @recursive
             def visit_ClassDef(self, node):
+                set_top_level_class = False
+
+                # Push the node on stack when there is a module level class
+                if not self.stack:
+                    set_top_level_class = True
+                    self.stack.append(node.name)
+
                 qualname = '.'.join(self.stack)
                 if self.qualname == qualname:
                     # decrement by one since lines starts with indexing by zero
                     self.line_number = node.lineno - 1
 
+                for child in ast.iter_child_nodes(node):
+                    if isinstance(child, ast.ClassDef):
+                        self.stack.append(child.name)
+                        self.visit(child)
+                        self.stack.pop()
+
+                # pop from stack here only when module level class is pushed
+                # This avoids the case where a class has children and we don't
+                # pop the child here which has to be done in the loop where it's pushed
+                if set_top_level_class and self.stack:
+                    self.stack.pop()
+
         qualname = object.__qualname__
         source = ''.join(lines)
         tree = ast.parse(source)
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index 24777a57fd21b2..bd7ddb16567165 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -201,18 +201,26 @@ class cls197:
     class cls201:
         pass
 
-#line 204
+class cls204:
+    class cls205:
+        class cls206:
+            pass
+    class cls208:
+        class cls209:
+            pass
+
+#line 212
 def positional_only_arg(a, /):
     pass
 
-#line 208
+#line 216
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 212
+# line 220
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 216
+#line 224
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 82b256492c2f33..f131f4f5fd1315 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -726,6 +726,13 @@ def test_class_decorator(self):
         self.assertSourceEqual(mod2.cls197, 195, 202)
         self.assertSourceEqual(mod2.cls197.cls201, 199, 202)
 
+    def test_multiple_children_classes(self):
+        self.assertSourceEqual(mod2.cls204, 204, 210)
+        self.assertSourceEqual(mod2.cls204.cls205, 205, 207)
+        self.assertSourceEqual(mod2.cls204.cls205.cls206, 206, 207)
+        self.assertSourceEqual(mod2.cls204.cls208, 208, 210)
+        self.assertSourceEqual(mod2.cls204.cls208.cls209, 209, 210)
+
 class TestNoEOL(GetSourceBase):
     def setUp(self):
         self.tempdir = TESTFN + '_dir'

From 22b9c9c94e08527adc5843a8dd7a3552f6ea695b Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 18:32:07 +0530
Subject: [PATCH 04/17] Fix PR comments

1. Use node.decorator_list to select decorators
2. Remove unwanted variables in ClassVisitor
3. Simplify stack management as per review
---
 Lib/inspect.py | 33 +++++++++------------------------
 1 file changed, 9 insertions(+), 24 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index b1705a5284d3ee..223e1ec1e591ac 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -806,52 +806,37 @@ def findsource(object):
 
         class ClassVisitor(ast.NodeVisitor):
 
-            def __init__(self, source, qualname, *args, **kwargs):
+            def __init__(self, qualname):
                 self.stack = []
-                self.source = source
                 self.line_number = None
                 self.qualname = qualname
-                super().__init__(*args, **kwargs)
 
             def visit_ClassDef(self, node):
-                set_top_level_class = False
-
-                # Push the node on stack when there is a module level class
-                if not self.stack:
-                    set_top_level_class = True
-                    self.stack.append(node.name)
+                self.stack.append(node.name)
 
                 qualname = '.'.join(self.stack)
                 if self.qualname == qualname:
+                    if node.decorator_list:
+                        line_number = node.decorator_list[0].lineno
+                    else:
+                        line_number = node.lineno
                     # decrement by one since lines starts with indexing by zero
-                    self.line_number = node.lineno - 1
+                    self.line_number = line_number - 1
 
                 for child in ast.iter_child_nodes(node):
                     if isinstance(child, ast.ClassDef):
-                        self.stack.append(child.name)
                         self.visit(child)
-                        self.stack.pop()
 
-                # pop from stack here only when module level class is pushed
-                # This avoids the case where a class has children and we don't
-                # pop the child here which has to be done in the loop where it's pushed
-                if set_top_level_class and self.stack:
-                    self.stack.pop()
+                self.stack.pop()
 
         qualname = object.__qualname__
         source = ''.join(lines)
         tree = ast.parse(source)
-        class_visitor = ClassVisitor(source, qualname)
+        class_visitor = ClassVisitor(qualname)
         class_visitor.visit(tree)
 
         if class_visitor.line_number is not None:
             line_number = class_visitor.line_number
-            decorator_pattern = re.compile(r'^(\s*@)')
-            for line in reversed(lines[:line_number]):
-                if decorator_pattern.match(line):
-                    line_number -= 1
-                else:
-                    break
             return lines, line_number
         else:
             raise OSError('could not find class definition')

From 2655b1d633d9bec5655080235fa0a802eefcb503 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 22:19:18 +0530
Subject: [PATCH 05/17] Add test for nested functions and async calls

---
 Lib/inspect.py              | 26 ++++++++++++------
 Lib/test/inspect_fodder2.py | 53 ++++++++++++++++++++++++++++---------
 Lib/test/test_inspect.py    | 26 +++++++++++++-----
 3 files changed, 77 insertions(+), 28 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 223e1ec1e591ac..52cb178eb0bc03 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -802,6 +802,8 @@ def findsource(object):
         return lines, 0
 
     if isclass(object):
+        # Lazy import ast because it's relatively heavy and
+        # it's not used for other than this part.
         import ast
 
         class ClassVisitor(ast.NodeVisitor):
@@ -811,22 +813,30 @@ def __init__(self, qualname):
                 self.line_number = None
                 self.qualname = qualname
 
-            def visit_ClassDef(self, node):
+            def visit_FunctionDef(self, node):
                 self.stack.append(node.name)
+                self.stack.append('<locals>')
+                self.generic_visit(node)
+                self.stack.pop()
+                self.stack.pop()
 
-                qualname = '.'.join(self.stack)
-                if self.qualname == qualname:
+            def visit_AsyncFunctionDef(self, node):
+                self.stack.append(node.name)
+                self.stack.append('<locals>')
+                self.generic_visit(node)
+                self.stack.pop()
+                self.stack.pop()
+
+            def visit_ClassDef(self, node):
+                self.stack.append(node.name)
+                if self.qualname == '.'.join(self.stack):
                     if node.decorator_list:
                         line_number = node.decorator_list[0].lineno
                     else:
                         line_number = node.lineno
                     # decrement by one since lines starts with indexing by zero
                     self.line_number = line_number - 1
-
-                for child in ast.iter_child_nodes(node):
-                    if isinstance(child, ast.ClassDef):
-                        self.visit(child)
-
+                self.generic_visit(node)
                 self.stack.pop()
 
         qualname = object.__qualname__
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index bd7ddb16567165..b2a7bd1027a925 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -187,40 +187,67 @@ class cls185:
         def func186(self):
             pass
 
-
 def class_decorator(cls):
     return cls
 
-# line 195
+# line 193
 @class_decorator
 @class_decorator
-class cls197:
+class cls196:
 
     @class_decorator
     @class_decorator
-    class cls201:
+    class cls200:
         pass
 
-class cls204:
-    class cls205:
-        class cls206:
+class cls203:
+    class cls204:
+        class cls205:
+            pass
+    class cls207:
+        class cls208:
             pass
-    class cls208:
-        class cls209:
+
+# line 211
+def func212():
+    class cls213:
+        def func(self):
+            pass
+    return cls213
+
+# line 218
+class cls213:
+    def func220(self):
+        class cls221:
+            pass
+        return cls221
+
+# line 225
+async def func226():
+    class cls227:
+        def func(self):
+            pass
+    return cls227
+
+# line 232
+class cls227:
+    async def func234(self):
+        class cls235:
             pass
+        return cls235
 
-#line 212
+#line 239
 def positional_only_arg(a, /):
     pass
 
-#line 216
+#line 243
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 220
+# line 247
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 224
+#line 251
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index f131f4f5fd1315..a597e11b93c4e8 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -473,6 +473,7 @@ def test_cleandoc(self):
     def test_getcomments(self):
         self.assertEqual(inspect.getcomments(mod), '# line 1\n')
         self.assertEqual(inspect.getcomments(mod.StupidGit), '# line 20\n')
+        self.assertEqual(inspect.getcomments(mod2.cls160), '# line 159\n')
         # If the object source file is not available, return None.
         co = compile('x=1', '_non_existing_filename.py', 'exec')
         self.assertIsNone(inspect.getcomments(co))
@@ -723,15 +724,26 @@ def test_nested_class_definition(self):
         self.assertSourceEqual(mod2.cls183.cls185, 185, 188)
 
     def test_class_decorator(self):
-        self.assertSourceEqual(mod2.cls197, 195, 202)
-        self.assertSourceEqual(mod2.cls197.cls201, 199, 202)
+        self.assertSourceEqual(mod2.cls196, 194, 201)
+        self.assertSourceEqual(mod2.cls196.cls200, 198, 201)
 
     def test_multiple_children_classes(self):
-        self.assertSourceEqual(mod2.cls204, 204, 210)
-        self.assertSourceEqual(mod2.cls204.cls205, 205, 207)
-        self.assertSourceEqual(mod2.cls204.cls205.cls206, 206, 207)
-        self.assertSourceEqual(mod2.cls204.cls208, 208, 210)
-        self.assertSourceEqual(mod2.cls204.cls208.cls209, 209, 210)
+        self.assertSourceEqual(mod2.cls203, 203, 209)
+        self.assertSourceEqual(mod2.cls203.cls204, 204, 206)
+        self.assertSourceEqual(mod2.cls203.cls204.cls205, 205, 206)
+        self.assertSourceEqual(mod2.cls203.cls207, 207, 209)
+        self.assertSourceEqual(mod2.cls203.cls207.cls208, 208, 209)
+
+    def test_nested_class_definition_inside_function(self):
+        self.assertSourceEqual(mod2.func212(), 213, 215)
+        self.assertSourceEqual(mod2.cls213, 219, 223)
+        self.assertSourceEqual(mod2.cls213().func220(), 221, 222)
+
+    def test_nested_class_definition_inside_async_function(self):
+        import asyncio
+        self.assertSourceEqual(asyncio.run(mod2.func226()), 227, 229)
+        self.assertSourceEqual(mod2.cls227, 233, 237)
+        self.assertSourceEqual(asyncio.run(mod2.cls227().func234()), 235, 236)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From bd51f498f5663af7906a744a8a38cd2a15d1fa59 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 22:23:42 +0530
Subject: [PATCH 06/17] Fix pydoc test since comments are returned now
 correctly

---
 Lib/test/test_pydoc.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Lib/test/test_pydoc.py b/Lib/test/test_pydoc.py
index 6d358f4fe2fc3c..ffabb7f1b94072 100644
--- a/Lib/test/test_pydoc.py
+++ b/Lib/test/test_pydoc.py
@@ -476,6 +476,7 @@ def test_getpager_with_stdin_none(self):
     def test_non_str_name(self):
         # issue14638
         # Treat illegal (non-str) name like no name
+
         class A:
             __name__ = 42
         class B:

From 558895fbe5f59183f575415c27745c8e732c2eaa Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 22:55:11 +0530
Subject: [PATCH 07/17] Set event loop policy as None to fix environment
 related change

---
 Lib/test/test_inspect.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index a597e11b93c4e8..6212fba00038fa 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -744,6 +744,7 @@ def test_nested_class_definition_inside_async_function(self):
         self.assertSourceEqual(asyncio.run(mod2.func226()), 227, 229)
         self.assertSourceEqual(mod2.cls227, 233, 237)
         self.assertSourceEqual(asyncio.run(mod2.cls227().func234()), 235, 236)
+        asyncio.set_event_loop_policy(None)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From 7d99c193ab7f16b23f7ee83c3a054d626b7ad3b7 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 23:37:31 +0530
Subject: [PATCH 08/17] Refactor visit_AsyncFunctionDef and tests

---
 Lib/inspect.py              |  7 +------
 Lib/test/inspect_fodder2.py | 40 ++++++++++++++++++-------------------
 Lib/test/test_inspect.py    | 14 ++++++-------
 3 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 52cb178eb0bc03..a139a05ddeed95 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -820,12 +820,7 @@ def visit_FunctionDef(self, node):
                 self.stack.pop()
                 self.stack.pop()
 
-            def visit_AsyncFunctionDef(self, node):
-                self.stack.append(node.name)
-                self.stack.append('<locals>')
-                self.generic_visit(node)
-                self.stack.pop()
-                self.stack.pop()
+            visit_AsyncFunctionDef = visit_FunctionDef
 
             def visit_ClassDef(self, node):
                 self.stack.append(node.name)
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index b2a7bd1027a925..91e46e362a13d7 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -205,49 +205,47 @@ class cls204:
         class cls205:
             pass
     class cls207:
-        class cls208:
+        class cls205:
             pass
 
 # line 211
 def func212():
     class cls213:
-        def func(self):
-            pass
+        pass
     return cls213
 
-# line 218
+# line 217
 class cls213:
-    def func220(self):
-        class cls221:
+    def func219(self):
+        class cls220:
             pass
-        return cls221
+        return cls220
 
-# line 225
-async def func226():
-    class cls227:
-        def func(self):
-            pass
-    return cls227
+# line 224
+async def func225():
+    class cls226:
+        pass
+    return cls226
 
-# line 232
-class cls227:
-    async def func234(self):
-        class cls235:
+# line 230
+class cls226:
+    async def func232(self):
+        class cls233:
             pass
         return cls235
 
-#line 239
+#line 237
 def positional_only_arg(a, /):
     pass
 
-#line 243
+#line 241
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 247
+# line 245
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 251
+#line 249
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 6212fba00038fa..b230a64f557438 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -732,19 +732,19 @@ def test_multiple_children_classes(self):
         self.assertSourceEqual(mod2.cls203.cls204, 204, 206)
         self.assertSourceEqual(mod2.cls203.cls204.cls205, 205, 206)
         self.assertSourceEqual(mod2.cls203.cls207, 207, 209)
-        self.assertSourceEqual(mod2.cls203.cls207.cls208, 208, 209)
+        self.assertSourceEqual(mod2.cls203.cls207.cls205, 208, 209)
 
     def test_nested_class_definition_inside_function(self):
         self.assertSourceEqual(mod2.func212(), 213, 215)
-        self.assertSourceEqual(mod2.cls213, 219, 223)
-        self.assertSourceEqual(mod2.cls213().func220(), 221, 222)
+        self.assertSourceEqual(mod2.cls213, 218, 222)
+        self.assertSourceEqual(mod2.cls213().func219(), 220, 221)
 
     def test_nested_class_definition_inside_async_function(self):
         import asyncio
-        self.assertSourceEqual(asyncio.run(mod2.func226()), 227, 229)
-        self.assertSourceEqual(mod2.cls227, 233, 237)
-        self.assertSourceEqual(asyncio.run(mod2.cls227().func234()), 235, 236)
-        asyncio.set_event_loop_policy(None)
+        self.assertSourceEqual(asyncio.run(mod2.func225()), 226, 228)
+        self.assertSourceEqual(mod2.cls226, 231, 235)
+        self.assertSourceEqual(asyncio.run(mod2.cls226().func232()), 233, 234)
+        self.addCleanup(asyncio.set_event_loop_policy, None)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From 5f843442716b111b3012b20094e794398311a4f3 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 23:58:29 +0530
Subject: [PATCH 09/17] Refactor to use local variables and fix tests

---
 Lib/inspect.py           | 30 ++++++++++++++----------------
 Lib/test/test_inspect.py |  6 +++---
 2 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index a139a05ddeed95..9dbf17114f1095 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -808,40 +808,38 @@ def findsource(object):
 
         class ClassVisitor(ast.NodeVisitor):
 
-            def __init__(self, qualname):
-                self.stack = []
-                self.line_number = None
-                self.qualname = qualname
-
             def visit_FunctionDef(self, node):
-                self.stack.append(node.name)
-                self.stack.append('<locals>')
+                stack.append(node.name)
+                stack.append('<locals>')
                 self.generic_visit(node)
-                self.stack.pop()
-                self.stack.pop()
+                stack.pop()
+                stack.pop()
 
             visit_AsyncFunctionDef = visit_FunctionDef
 
             def visit_ClassDef(self, node):
-                self.stack.append(node.name)
-                if self.qualname == '.'.join(self.stack):
+                nonlocal line_number
+                stack.append(node.name)
+                if qualname == '.'.join(stack):
                     if node.decorator_list:
                         line_number = node.decorator_list[0].lineno
                     else:
                         line_number = node.lineno
                     # decrement by one since lines starts with indexing by zero
-                    self.line_number = line_number - 1
+                    line_number = line_number - 1
                 self.generic_visit(node)
-                self.stack.pop()
+                stack.pop()
 
+        stack = []
+        line_number = None
         qualname = object.__qualname__
         source = ''.join(lines)
         tree = ast.parse(source)
-        class_visitor = ClassVisitor(qualname)
+        class_visitor = ClassVisitor()
         class_visitor.visit(tree)
 
-        if class_visitor.line_number is not None:
-            line_number = class_visitor.line_number
+        if line_number is not None:
+            line_number = line_number
             return lines, line_number
         else:
             raise OSError('could not find class definition')
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index b230a64f557438..44cfa8999fd0aa 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -735,16 +735,16 @@ def test_multiple_children_classes(self):
         self.assertSourceEqual(mod2.cls203.cls207.cls205, 208, 209)
 
     def test_nested_class_definition_inside_function(self):
-        self.assertSourceEqual(mod2.func212(), 213, 215)
+        self.assertSourceEqual(mod2.func212(), 213, 214)
         self.assertSourceEqual(mod2.cls213, 218, 222)
         self.assertSourceEqual(mod2.cls213().func219(), 220, 221)
 
     def test_nested_class_definition_inside_async_function(self):
         import asyncio
-        self.assertSourceEqual(asyncio.run(mod2.func225()), 226, 228)
+        self.assertSourceEqual(asyncio.run(mod2.func225()), 226, 227)
         self.assertSourceEqual(mod2.cls226, 231, 235)
         self.assertSourceEqual(asyncio.run(mod2.cls226().func232()), 233, 234)
-        self.addCleanup(asyncio.set_event_loop_policy, None)
+        asyncio.set_event_loop_policy(None)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From 8a63f36d38269b7972c9c8397e400e0d0be282f5 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sat, 3 Nov 2018 23:59:09 +0530
Subject: [PATCH 10/17] Add patch attribution

---
 .../next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst       | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
index 8919ad663493b9..df113f821af667 100644
--- a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
+++ b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
@@ -1,3 +1,3 @@
 :meth:`inspect.getsource` now returns correct source code for inner class
 with same name as module level class. class decorators are also returned as
-part of the source.
+part of the source. Patch by Karthikeyan Singaravelan.

From 4c7a61b3c476b21b6e6c9f0d3ccfe4f17042495a Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sun, 4 Nov 2018 00:02:38 +0530
Subject: [PATCH 11/17] Use self.addCleanup for asyncio

---
 Lib/test/test_inspect.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 44cfa8999fd0aa..6d81db58398b36 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -744,7 +744,7 @@ def test_nested_class_definition_inside_async_function(self):
         self.assertSourceEqual(asyncio.run(mod2.func225()), 226, 227)
         self.assertSourceEqual(mod2.cls226, 231, 235)
         self.assertSourceEqual(asyncio.run(mod2.cls226().func232()), 233, 234)
-        asyncio.set_event_loop_policy(None)
+        self.addCleanup(asyncio.set_event_loop_policy, None)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):

From 30865f60b45cef353c1936225718347e1fca9f88 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Sun, 4 Nov 2018 00:30:27 +0530
Subject: [PATCH 12/17] Rename ClassVisitor to ClassFinder and fix asyncio
 cleanup

---
 Lib/inspect.py                                             | 7 +++----
 Lib/test/test_inspect.py                                   | 2 +-
 .../next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst  | 2 +-
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 9dbf17114f1095..3e30be39984851 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -806,7 +806,7 @@ def findsource(object):
         # it's not used for other than this part.
         import ast
 
-        class ClassVisitor(ast.NodeVisitor):
+        class ClassFinder(ast.NodeVisitor):
 
             def visit_FunctionDef(self, node):
                 stack.append(node.name)
@@ -835,11 +835,10 @@ def visit_ClassDef(self, node):
         qualname = object.__qualname__
         source = ''.join(lines)
         tree = ast.parse(source)
-        class_visitor = ClassVisitor()
-        class_visitor.visit(tree)
+        class_finder = ClassFinder()
+        class_finder.visit(tree)
 
         if line_number is not None:
-            line_number = line_number
             return lines, line_number
         else:
             raise OSError('could not find class definition')
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 6d81db58398b36..420e8bf448973e 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -741,10 +741,10 @@ def test_nested_class_definition_inside_function(self):
 
     def test_nested_class_definition_inside_async_function(self):
         import asyncio
+        self.addCleanup(asyncio.set_event_loop_policy, None)
         self.assertSourceEqual(asyncio.run(mod2.func225()), 226, 227)
         self.assertSourceEqual(mod2.cls226, 231, 235)
         self.assertSourceEqual(asyncio.run(mod2.cls226().func232()), 233, 234)
-        self.addCleanup(asyncio.set_event_loop_policy, None)
 
 class TestNoEOL(GetSourceBase):
     def setUp(self):
diff --git a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
index df113f821af667..00385c87977f05 100644
--- a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
+++ b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
@@ -1,3 +1,3 @@
 :meth:`inspect.getsource` now returns correct source code for inner class
-with same name as module level class. class decorators are also returned as
+with same name as module level class. Class decorators are also returned as
 part of the source. Patch by Karthikeyan Singaravelan.

From 054c3179ea438456f1490baf541d44f01d509151 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Fri, 25 Jan 2019 10:06:49 +0530
Subject: [PATCH 13/17] Return first class inside conditional in case of
 multiple definitions. Remove decorator for class source.

---
 Lib/inspect.py              | 14 ++++++--------
 Lib/test/inspect_fodder2.py | 19 ++++++++++++++-----
 Lib/test/test_inspect.py    |  8 ++++++--
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 3e30be39984851..2b47e94521b368 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -821,12 +821,9 @@ def visit_ClassDef(self, node):
                 nonlocal line_number
                 stack.append(node.name)
                 if qualname == '.'.join(stack):
-                    if node.decorator_list:
-                        line_number = node.decorator_list[0].lineno
-                    else:
-                        line_number = node.lineno
                     # decrement by one since lines starts with indexing by zero
-                    line_number = line_number - 1
+                    line_number = node.lineno - 1
+                    raise StopIteration(line_number)
                 self.generic_visit(node)
                 stack.pop()
 
@@ -836,9 +833,10 @@ def visit_ClassDef(self, node):
         source = ''.join(lines)
         tree = ast.parse(source)
         class_finder = ClassFinder()
-        class_finder.visit(tree)
-
-        if line_number is not None:
+        try:
+            class_finder.visit(tree)
+        except StopIteration as e:
+            line_number = e.value
             return lines, line_number
         else:
             raise OSError('could not find class definition')
diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index 91e46e362a13d7..040a50f597c722 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -232,20 +232,29 @@ class cls226:
     async def func232(self):
         class cls233:
             pass
-        return cls235
+        return cls233
 
-#line 237
+if True:
+    class cls238:
+        class cls239:
+            pass
+else:
+    class cls238:
+        class cls239:
+            pass
+
+#line 246
 def positional_only_arg(a, /):
     pass
 
-#line 241
+#line 250
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 245
+# line 254
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 249
+#line 258
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 420e8bf448973e..682daa38b7533b 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -724,8 +724,12 @@ def test_nested_class_definition(self):
         self.assertSourceEqual(mod2.cls183.cls185, 185, 188)
 
     def test_class_decorator(self):
-        self.assertSourceEqual(mod2.cls196, 194, 201)
-        self.assertSourceEqual(mod2.cls196.cls200, 198, 201)
+        self.assertSourceEqual(mod2.cls196, 196, 201)
+        self.assertSourceEqual(mod2.cls196.cls200, 200, 201)
+
+    def test_class_inside_conditional(self):
+        self.assertSourceEqual(mod2.cls238, 238, 240)
+        self.assertSourceEqual(mod2.cls238.cls239, 239, 240)
 
     def test_multiple_children_classes(self):
         self.assertSourceEqual(mod2.cls203, 203, 209)

From 26d1c9616dd380b4bfd41751178e2cfada186d18 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Fri, 25 Jan 2019 10:23:41 +0530
Subject: [PATCH 14/17] Add docstring to make the test correct

---
 Lib/test/inspect_fodder2.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/Lib/test/inspect_fodder2.py b/Lib/test/inspect_fodder2.py
index 040a50f597c722..e7d4b53ebefcc6 100644
--- a/Lib/test/inspect_fodder2.py
+++ b/Lib/test/inspect_fodder2.py
@@ -237,24 +237,25 @@ class cls233:
 if True:
     class cls238:
         class cls239:
-            pass
+            '''if clause cls239'''
 else:
     class cls238:
         class cls239:
+            '''else clause 239'''
             pass
 
-#line 246
+#line 247
 def positional_only_arg(a, /):
     pass
 
-#line 250
+#line 251
 def all_markers(a, b, /, c, d, *, e, f):
     pass
 
-# line 254
+# line 255
 def all_markers_with_args_and_kwargs(a, b, /, c, d, *args, e, f, **kwargs):
     pass
 
-#line 258
+#line 259
 def all_markers_with_defaults(a, b=1, /, c=2, d=3, *, e=4, f=5):
     pass

From d7ae710457be3e8d5c276ef775245b242337d2a3 Mon Sep 17 00:00:00 2001
From: xtreak <tir.karthi@gmail.com>
Date: Fri, 25 Jan 2019 10:27:19 +0530
Subject: [PATCH 15/17] Modify NEWS entry regarding decorators

---
 .../next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst      | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
index 00385c87977f05..ffbd88ce3fb9e2 100644
--- a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
+++ b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
@@ -1,3 +1,2 @@
 :meth:`inspect.getsource` now returns correct source code for inner class
-with same name as module level class. Class decorators are also returned as
-part of the source. Patch by Karthikeyan Singaravelan.
+with same name as module level class. Patch by Karthikeyan Singaravelan.

From aef42284a16b477535534f827992b327d9a202b9 Mon Sep 17 00:00:00 2001
From: Karthikeyan Singaravelan <tir.karthi@gmail.com>
Date: Wed, 1 May 2019 22:42:40 +0530
Subject: [PATCH 16/17] Return decorators too for bpo-15856

---
 Lib/inspect.py                                            | 8 +++++++-
 Lib/test/test_inspect.py                                  | 4 ++--
 .../next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst | 3 ++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index 2b47e94521b368..ec0d791fe3197d 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -821,8 +821,14 @@ def visit_ClassDef(self, node):
                 nonlocal line_number
                 stack.append(node.name)
                 if qualname == '.'.join(stack):
+                    # Return the decorator for the class if present
+                    if node.decorator_list:
+                        line_number = node.decorator_list[0].lineno
+                    else:
+                        line_number = node.lineno
+
                     # decrement by one since lines starts with indexing by zero
-                    line_number = node.lineno - 1
+                    line_number -= 1
                     raise StopIteration(line_number)
                 self.generic_visit(node)
                 stack.pop()
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 682daa38b7533b..98a9c0a662a093 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -724,8 +724,8 @@ def test_nested_class_definition(self):
         self.assertSourceEqual(mod2.cls183.cls185, 185, 188)
 
     def test_class_decorator(self):
-        self.assertSourceEqual(mod2.cls196, 196, 201)
-        self.assertSourceEqual(mod2.cls196.cls200, 200, 201)
+        self.assertSourceEqual(mod2.cls196, 194, 201)
+        self.assertSourceEqual(mod2.cls196.cls200, 198, 201)
 
     def test_class_inside_conditional(self):
         self.assertSourceEqual(mod2.cls238, 238, 240)
diff --git a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
index ffbd88ce3fb9e2..bf6b672964fa64 100644
--- a/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
+++ b/Misc/NEWS.d/next/Library/2018-11-03-16-18-20.bpo-35113.vwvWKG.rst
@@ -1,2 +1,3 @@
 :meth:`inspect.getsource` now returns correct source code for inner class
-with same name as module level class. Patch by Karthikeyan Singaravelan.
+with same name as module level class. Decorators are also returned as part
+of source of the class. Patch by Karthikeyan Singaravelan.

From 797c48dd82e37849ea14e7fb907128df49b06ac7 Mon Sep 17 00:00:00 2001
From: Karthikeyan Singaravelan <tir.karthi@gmail.com>
Date: Wed, 1 May 2019 23:30:14 +0530
Subject: [PATCH 17/17] Move ast and the class source code to top. Use proper
 Exception.

---
 Lib/inspect.py | 76 ++++++++++++++++++++++++++------------------------
 1 file changed, 40 insertions(+), 36 deletions(-)

diff --git a/Lib/inspect.py b/Lib/inspect.py
index ec0d791fe3197d..ad7e8cb1203e7f 100644
--- a/Lib/inspect.py
+++ b/Lib/inspect.py
@@ -32,6 +32,7 @@
               'Yury Selivanov <yselivanov@sprymix.com>')
 
 import abc
+import ast
 import dis
 import collections.abc
 import enum
@@ -770,6 +771,42 @@ def getmodule(object, _filename=None):
         if builtinobject is object:
             return builtin
 
+
+class ClassFoundException(Exception):
+    pass
+
+
+class _ClassFinder(ast.NodeVisitor):
+
+    def __init__(self, qualname):
+        self.stack = []
+        self.qualname = qualname
+
+    def visit_FunctionDef(self, node):
+        self.stack.append(node.name)
+        self.stack.append('<locals>')
+        self.generic_visit(node)
+        self.stack.pop()
+        self.stack.pop()
+
+    visit_AsyncFunctionDef = visit_FunctionDef
+
+    def visit_ClassDef(self, node):
+        self.stack.append(node.name)
+        if self.qualname == '.'.join(self.stack):
+            # Return the decorator for the class if present
+            if node.decorator_list:
+                line_number = node.decorator_list[0].lineno
+            else:
+                line_number = node.lineno
+
+            # decrement by one since lines starts with indexing by zero
+            line_number -= 1
+            raise ClassFoundException(line_number)
+        self.generic_visit(node)
+        self.stack.pop()
+
+
 def findsource(object):
     """Return the entire source file and starting line number for an object.
 
@@ -802,47 +839,14 @@ def findsource(object):
         return lines, 0
 
     if isclass(object):
-        # Lazy import ast because it's relatively heavy and
-        # it's not used for other than this part.
-        import ast
-
-        class ClassFinder(ast.NodeVisitor):
-
-            def visit_FunctionDef(self, node):
-                stack.append(node.name)
-                stack.append('<locals>')
-                self.generic_visit(node)
-                stack.pop()
-                stack.pop()
-
-            visit_AsyncFunctionDef = visit_FunctionDef
-
-            def visit_ClassDef(self, node):
-                nonlocal line_number
-                stack.append(node.name)
-                if qualname == '.'.join(stack):
-                    # Return the decorator for the class if present
-                    if node.decorator_list:
-                        line_number = node.decorator_list[0].lineno
-                    else:
-                        line_number = node.lineno
-
-                    # decrement by one since lines starts with indexing by zero
-                    line_number -= 1
-                    raise StopIteration(line_number)
-                self.generic_visit(node)
-                stack.pop()
-
-        stack = []
-        line_number = None
         qualname = object.__qualname__
         source = ''.join(lines)
         tree = ast.parse(source)
-        class_finder = ClassFinder()
+        class_finder = _ClassFinder(qualname)
         try:
             class_finder.visit(tree)
-        except StopIteration as e:
-            line_number = e.value
+        except ClassFoundException as e:
+            line_number = e.args[0]
             return lines, line_number
         else:
             raise OSError('could not find class definition')