Added Import Changes

Author: hanse7962

src/inline/plugin.py

@@ -296,11 +296,17 @@ class ExtractInlineTest(ast.NodeTransformer):
     arg_timeout_str = "timeout"
 
     assume = "assume"
+    inline_module_imported = False
+    import_str = "import"
+    from_str = "from"
+    as_str = "as"
+    
     inline_module_imported = False
 
     def __init__(self):
         self.cur_inline_test = InlineTest()
         self.inline_test_list = []
+        self.import_list = []
 
     def is_inline_test_class(self, node):
         if isinstance(node, ast.Call):
@@ -350,16 +356,39 @@ def find_previous_stmt(self, node):
                     return prev_stmt_node
             return self.find_condition_stmt(prev_stmt_node)
 
-    def collect_inline_test_calls(self, node, inline_test_calls: List[ast.Call]):
+    def collect_inline_test_calls(self, node, inline_test_calls: List[ast.Call], import_calls: List[ast.Import], import_from_calls: List[ast.ImportFrom]):
         """
         collect all function calls in the node
         """
         if isinstance(node, ast.Attribute):
-            self.collect_inline_test_calls(node.value, inline_test_calls)
+            self.collect_inline_test_calls(node.value, inline_test_calls, import_calls, import_from_calls)
         elif isinstance(node, ast.Call):
             inline_test_calls.append(node)
-            self.collect_inline_test_calls(node.func, inline_test_calls)
+            self.collect_inline_test_calls(node.func, inline_test_calls, import_calls, import_from_calls)
+        elif isinstance(node, ast.Import):
+            import_calls.append(node)
+            self.collect_inline_test_calls(node.func, inline_test_calls, import_calls, import_from_calls)
+        elif isinstance(node, ast.ImportFrom):
+            import_from_calls.append(node)
+            self.collect_inline_test_calls(node.func, inline_test_calls, import_calls, import_from_calls)
+
+    def collect_import_calls(self, node, import_calls: List[ast.Import], import_from_calls: List[ast.ImportFrom]):
+        """
+        collect all import calls in the node (should be done first)
+        """
 
+        while not isinstance(node, ast.Module) and node.parent != None:
+            node = node.parent
+       
+        if not isinstance(node, ast.Module):
+            return
+            
+        for child in node.children:
+            if isinstance(child, ast.Import):
+                import_calls.append(child)
+            elif isinstance(child, ast.ImportFrom):
+                import_from_calls.append(child)
+                
     def parse_constructor(self, node):
         """
         Parse a constructor call.
@@ -931,8 +960,13 @@ def parse_parameterized_test(self):
             parameterized_test.test_name = self.cur_inline_test.test_name + "_" + str(index)
 
     def parse_inline_test(self, node):
-        inline_test_calls = []
-        self.collect_inline_test_calls(node, inline_test_calls)
+        import_calls = []
+        import_from_calls = []
+        inline_test_calls = [] 
+        
+        self.collect_inline_test_calls(node, inline_test_calls, import_calls, import_from_calls)
+        self.collect_import_calls(node, import_calls, import_from_calls)
+        
         inline_test_calls.reverse()
 
         if len(inline_test_calls) <= 1:
@@ -953,14 +987,32 @@ def parse_inline_test(self, node):
                 self.parse_assume(call)
                 inline_test_call_index += 1
 
-        # "given(a, 1)"
         for call in inline_test_calls[inline_test_call_index:]:
-            if isinstance(call.func, ast.Attribute) and call.func.attr == self.given_str:
-                self.parse_given(call)
-                inline_test_call_index += 1
+            if isinstance(call.func, ast.Attribute):
+                if call.func.attr == self.given_str:
+                    self.parse_given(call)
+                    inline_test_call_index += 1
+                elif call.func.attr == self.diff_given_str:
+                    self.parse_diff_given(call)
+                    inline_test_call_index += 1
+                
+                # match call.func.attr:
+                #     # "given(a, 1)"
+                #     case self.given_str:
+                #         self.parse_given(call)
+                #         inline_test_call_index += 1
+                #      # "diff_given(devices, ["cpu", "cuda"])"
+                #     case self.diff_given_str:
+                #         self.parse_diff_given(call)
+                #         inline_test_call_index += 1
             else:
                 break
 
+        for import_stmt in import_calls:
+            self.cur_inline_test.import_stmts.append(import_stmt)
+        for import_stmt in import_from_calls:
+            self.cur_inline_test.import_stmts.append(import_stmt)
+
         # "check_eq" or "check_true" or "check_false" or "check_neq"
         for call in inline_test_calls[inline_test_call_index:]:
             # "check_eq(a, 1)"

tests/test_plugin.py

@@ -2,9 +2,48 @@
 from _pytest.pytester import Pytester
 import pytest
 
+# For testing in Spyder only
+if __name__ == "__main__":
+    pytest.main(['-v', '-s']) 
+
 
 # pytest -p pytester
 class TestInlinetests:
+    def test_inline_detects_imports(self, pytester: Pytester):
+        checkfile = pytester.makepyfile(
+            """ 
+        from inline import itest
+        import datetime
+        
+        def m(a):
+            b = a + datetime.timedelta(days=365)
+            itest().given(a, datetime.timedelta(days=1)).check_eq(b, datetime.timedelta(days=366))
+    """
+        )
+        for x in (pytester.path, checkfile):
+            items, reprec = pytester.inline_genitems(x)
+            assert len(items) == 1
+            res = pytester.runpytest()
+            assert res.ret != 1
+
+    # def test_inline_detects_from_imports(self, pytester: Pytester):
+    #     checkfile = pytester.makepyfile(
+    #         """ 
+    #     from inline import itest
+    #     import numpy as np
+    #     from scipy import stats as st
+        
+    #     def m(n, p):
+    #         b = st.binom(n, p)
+    #         itest().given(n, 100).given(p, 0.5).check_eq(b.mean(), n * p)
+    # """
+    #     )
+    #     for x in (pytester.path, checkfile):
+    #         items, reprec = pytester.inline_genitems(x)
+    #         assert len(items) == 1
+    #         res = pytester.runpytest()
+    #         assert res.ret == 0
+
     def test_inline_parser(self, pytester: Pytester):
         checkfile = pytester.makepyfile(
             """ 
@@ -31,6 +70,7 @@ def m(a):
             items, reprec = pytester.inline_genitems(x)
             assert len(items) == 0
 
+
     def test_inline_malformed_given(self, pytester: Pytester):
         checkfile = pytester.makepyfile(
             """ 
@@ -118,21 +158,6 @@ def m(a):
             res = pytester.runpytest()
             assert res.ret == 0
 
-    def test_check_eq_parameterized_tests(self, pytester: Pytester):
-        checkfile = pytester.makepyfile(
-            """ 
-        from inline import itest
-        def m(a):
-            a = a + 1
-            itest(parameterized=True).given(a, [2, 3]).check_eq(a, [3, 4])
-    """
-        )
-        for x in (pytester.path, checkfile):
-            items, reprec = pytester.inline_genitems(x)
-            assert len(items) == 2
-            res = pytester.runpytest()
-            assert res.ret == 0
-
     def test_malformed_check_eq_parameterized_tests(self, pytester: Pytester):
         checkfile = pytester.makepyfile(
             """