Iter raise statements in a Function (#16)

* Add Function.iter_raise_stmts method and tests

* Add Alisdair Robertson to AUTHORS.txt

* Cleanup Function.iter_raise_stmts and test

Decided not to try and exclude exceptions that would be caught by a try-catch

AUTHORS.txt

@@ -5,6 +5,7 @@ David Halter (@davidhalter) <davidhalter88@gmail.com>
 
 Code Contributors
 =================
+Alisdair Robertson (@robodair)
 
 
 Code Contributors (to Jedi and therefore possibly to this library)

parso/python/tree.py

@@ -592,6 +592,21 @@ class Function(ClassOrFunc):
 
         return scan(self.children)
 
+    def iter_raise_stmts(self):
+        """
+        Returns a generator of `raise_stmt`. Includes raise statements inside try-except blocks
+        """
+        def scan(children):
+            for element in children:
+                if element.type == 'raise_stmt' \
+                        or element.type == 'keyword' and element.value == 'raise':
+                    yield element
+                if element.type in _RETURN_STMT_CONTAINERS:
+                    for e in scan(element.children):
+                        yield e
+
+        return scan(self.children)
+
     def is_generator(self):
         """
         :return bool: Checks if a function is a generator or not.

test/test_parser_tree.py

@@ -125,6 +125,10 @@ def get_return_stmts(code):
     return list(parse(code).children[0].iter_return_stmts())
 
 
+def get_raise_stmts(code, child):
+    return list(parse(code).children[child].iter_raise_stmts())
+
+
 def test_yields(each_version):
     y, = get_yield_exprs('def x(): yield', each_version)
     assert y.value == 'yield'
@@ -149,3 +153,30 @@ def test_returns():
 
     r, = get_return_stmts('def x(): return 1')
     assert r.type == 'return_stmt'
+
+
+def test_raises():
+    code = """
+def single_function():
+    raise Exception
+def top_function():
+    def inner_function():
+        raise NotImplementedError()
+    inner_function()
+    raise Exception
+def top_function_three():
+    try:
+        raise NotImplementedError()
+    except NotImplementedError:
+        pass
+    raise Exception
+    """
+
+    r = get_raise_stmts(code, 0) #  Lists in a simple Function
+    assert len(list(r)) == 1
+
+    r = get_raise_stmts(code, 1) #  Doesn't Exceptions list in closures
+    assert len(list(r)) == 1
+
+    r = get_raise_stmts(code, 2) #  Lists inside try-catch
+    assert len(list(r)) == 2