fix: statement variables in python

Author: SunBK201

scubatrace/python/statement.py

@@ -2,6 +2,7 @@
 
 from functools import cached_property
 
+from ..identifier import Identifier
 from ..statement import BlockStatement, SimpleStatement, Statement
 
 
@@ -31,6 +32,30 @@ def right_uncle_ancestor(self) -> Statement | None:
             cur = cur.parent
         return None
 
+    @cached_property
+    def variables(self) -> list[Identifier]:
+        """
+        Variables in the statement.
+        """
+        variables = []
+        for identifier in self.identifiers:
+            node = identifier.node
+            parent = node.parent
+            if parent is not None:
+                if parent.type in [
+                    "ERROR",
+                    "call",
+                    "function_definition",
+                ]:
+                    continue
+                if (
+                    parent.type == "keyword_argument"
+                    and parent.child_by_field_name("name") == node
+                ):
+                    continue
+            variables.append(identifier)
+        return variables
+
 
 class PythonBlockStatement(BlockStatement):
     @cached_property

tests/samples/python/test.py

@@ -1,11 +1,22 @@
 import scubatrace
 
 
+def add(a, b):
+    return a + b
+
+
+def main():
+    a = 3
+    b = 4
+    c = add(a, b=b)
+    return c
+
+
 def test_cfg():
     project = scubatrace.PythonProject(".")
     file = project.files["test.py"]
     file.export_cfg_dot("test.dot")
 
 
 if __name__ == "__main__":
-    test_cfg()
+    main()

tests/test_statement.py

@@ -65,6 +65,29 @@ def test_statement_is_taint_from_entry(self):
         self.assertIsInstance(is_taint_from_entry, bool)
 
 
+class TestPythonStatement(unittest.TestCase):
+    def setUp(self):
+        self.test_dir = Path(__file__).parent
+        self.samples_dir = self.test_dir / "samples"
+        self.project_path = self.samples_dir / "python"
+        self.project = scubatrace.Project.create(
+            str(self.project_path), language=scubatrace.language.PYTHON
+        )
+        self.file = self.project.files.get("test.py") or self.fail()
+        statement = self.file.statements_by_line(11)
+        self.statement = statement[0]
+
+    def test_statement_create(self):
+        statement = scubatrace.SimpleStatement.create(
+            self.statement.node, self.statement.parent
+        )
+        self.assertIsNotNone(statement)
+
+    def test_statement_variables(self):
+        variables = self.statement.variables
+        self.assertEqual(len(variables), 3)
+
+
 class TestJavaScriptStatement(unittest.TestCase):
     def setUp(self):
         self.test_dir = Path(__file__).parent
@@ -86,3 +109,13 @@ def test_statement_create(self):
     def test_statement_walk_backward(self):
         for stmt in self.statement.walk_backward(depth=3, base="control"):
             self.assertIn(stmt.start_line, [1, 2, 3, 4])
+
+
+if __name__ == "__main__":
+    project = scubatrace.Project.create(
+        "/Users/sunbk201/Downloads/python-xxhash",
+        language=scubatrace.language.PYTHON,
+        enable_lsp=True,
+    )
+    for impor in project.files["setup.py"].imports:
+        print(impor._path)