fix: Added support for more AST classes

stencila/pyla/code_parsing.py

@@ -351,9 +351,10 @@ def parse(self, chunk: CodeChunk) -> CodeChunkParseResult:
         return CodeChunkParseResult(chunk_ast, self.imports, self.assigns, self.declares, self.alters, self.uses,
                                     self.reads)
 
-    # pylint: disable=R0912  # Too many branches warning but this is kind of a special case
+    # pylint: disable=R0912,R0915 # Too many branches/statements warning but this is kind of a special case
     def parse_statement(self,
-                        statement: typing.Union[ast.stmt, ast.expr, typing.Sequence[typing.Union[ast.stmt, ast.expr]]]
+                        statement: typing.Union[
+                            ast.stmt, ast.expr, typing.Sequence[typing.Union[ast.stmt, ast.expr, ast.comprehension]]]
                         ) -> None:
         """General statement parser that delegates to parsers for specific parser types."""
         if isinstance(statement, list):
@@ -407,7 +408,8 @@ def parse_statement(self,
             self._parse_lambda(statement)
         elif isinstance(statement, ast.UnaryOp):
             self._parse_unary_op(statement)
-        elif isinstance(statement, (ast.ClassDef, ast.Num, ast.Str, ast.Pass, ast.NameConstant)):
+        elif isinstance(statement, (
+                ast.ClassDef, ast.Num, ast.Str, ast.Pass, ast.NameConstant, ast.Bytes, ast.Break, ast.Continue)):
             pass
         else:
             raise TypeError('Unrecognized statement: {}'.format(statement))
@@ -689,9 +691,9 @@ def _parse_comprehension(self, statement: ast.comprehension) -> None:
         target = statement.target
 
         if hasattr(target, 'elts'):
-            self.name_skip = [name.id for name in target.elts]
+            self.name_skip = [name.id for name in target.elts]  # type: ignore
         elif hasattr(target, 'id'):
-            self.name_skip = [target.id]
+            self.name_skip = [target.id]  # type: ignore
 
         self.parse_statement(statement.iter)
         self.parse_statement(statement.ifs)

tests/test_code_parsing.py

@@ -14,6 +14,7 @@
 c: int = 10
 def test_func():
     d = 4
+e = b'abc123'
 """
 
 FUNCTION_CODE = """
@@ -66,6 +67,7 @@ def basic():  # don't add it twice
 ff.gg
 hh[ii]
 jj.kk.ll
+mm % nn
 """
 
 OPEN_CODE = """
@@ -169,8 +171,12 @@ def open_func():
 FOR_CODE = """
 for a in range(10):
     print("{}".format(d))
+    break
 else:
     print(e)
+    
+for b in range(10):
+    continue
 """
 
 EXCEPT_CODE = """
@@ -227,7 +233,7 @@ def test_variable_parsing() -> None:
 
     assert type(parse_result.declares[1]) == Function  # The correctness of parsing the function is tested elsewhere
 
-    assert parse_result.assigns == ['a', 'b']
+    assert parse_result.assigns == ['a', 'b', 'e']
 
     check_result_fields_empty(parse_result, ['declares', 'assigns'])
 
@@ -295,7 +301,7 @@ def test_uses_parsing():
     check_result_fields_empty(parse_result, ['uses'])
 
     uses = ['a', 'b', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
-            'w', 'x', 'y', 'z', 'aa', 'bb', 'cc', 'dd', 'ee', 'ff', 'hh', 'ii', 'jj']
+            'w', 'x', 'y', 'z', 'aa', 'bb', 'cc', 'dd', 'ee', 'ff', 'hh', 'ii', 'jj', 'mm', 'nn']
 
     assert sorted(uses) == sorted(parse_result.uses)
 
@@ -373,7 +379,7 @@ def test_conditional_code_parsing():
 def test_for_parsing():
     parse_result = parse_code(FOR_CODE)
 
-    assert ['a'] == parse_result.assigns
+    assert ['a', 'b'] == parse_result.assigns
     assert ['d', 'e'] == sorted(parse_result.uses)
 
     check_result_fields_empty(parse_result, ['assigns', 'uses'])