Adding scopes "File", "Function", and "Class" for Python

src/cleanup_rules/python/scope_config.toml

@@ -0,0 +1,47 @@
+# Scope generator for python files
+[[scopes]]
+name = "Function"
+[[scopes.rules]]
+enclosing_node = """
+(
+  (function_definition
+      name: (_) @n
+      parameters: (parameters) @fp
+  ) @xdn
+)"""
+scope = """
+(
+  [
+    (function_definition
+        name: (_) @z
+        parameters: (parameters) @tp
+    ) @qdn
+    (#eq? @z "@n")
+    (#eq? @tp "@fp")
+  ]
+)
+"""
+
+[[scopes]]
+name = "Class"
+[[scopes.rules]]
+enclosing_node = """
+((class_definition name: (_) @n) @c)
+"""
+scope = """
+(
+  ((class_definition
+      name: (_) @z) @qc)
+  (#eq? @z "@n")
+)
+"""
+
+
+[[scopes]]
+name = "File"
+[[scopes.rules]]
+enclosing_node = """
+(module) @p_m
+"""
+scope = "(module) @python_module"
+

src/models/language.rs

@@ -198,7 +198,9 @@ impl std::str::FromStr for PiranhaLanguage {
         language: tree_sitter_python::language(),
         rules: None,
         edges: None,
-        scopes: vec![],
+        scopes: parse_toml::<ScopeConfig>(include_str!("../cleanup_rules/python/scope_config.toml"))
+            .scopes()
+            .to_vec(),
         comment_nodes: vec![],
       }),
       SWIFT => {

test-resources/python/find_import_extend_method_chain/configurations/edges.toml

@@ -0,0 +1,4 @@
+[[edges]]
+scope = "File"
+from = "find_import"
+to = ["extend_method_chain"]

test-resources/python/find_import_extend_method_chain/configurations/rules.toml

@@ -0,0 +1,59 @@
+[[rules]]
+name = "find_import"
+is_seed_rule = true
+query = """
+(
+  (import_from_statement
+    module_name: (dotted_name
+      (identifier) @module_name_part
+      (identifier) @submodule_name_part)
+    name: (dotted_name
+      (identifier) @imported_name)
+  (#eq? @module_name_part @module_name_to_check)
+  (#eq? @submodule_name_part @submodule_name_to_check)
+  (#eq? @imported_name @imported_name_to_check))
+)
+"""
+holes = [
+  "imported_name_to_check",
+  "submodule_name_to_check",
+  "module_name_to_check",
+]
+
+[[rules]]
+name = "extend_method_chain"
+query = """
+(
+  (call
+    function: (attribute
+      object: (attribute
+        object: (identifier) @obj
+        attribute: (identifier) @builder_method
+      )
+      attribute: (identifier) @config_method
+    )
+    arguments: (argument_list
+      (string) @config_arg1
+      (string) @config_arg2
+    )
+  )
+  (#eq? @obj @imported_name_to_check)
+  (#eq? @builder_method @class_attribute)
+  (#eq? @config_method @method_call_to_check)
+)
+"""
+replace_node = "config_method"
+replace = "@method_call_to_add.@config_method"
+holes = [
+  "imported_name_to_check",
+  "class_attribute",
+  "method_call_to_check",
+  "class_attribute",
+  "method_call_to_check",
+  "method_call_to_add",
+]
+is_seed_rule = false
+
+[[rules.filters]]
+enclosing_node = "((call) @cal)"
+not_contains = ["cs :[other].@method_call_to_add"]

test-resources/python/find_import_extend_method_chain/expected/method_chain.py

@@ -0,0 +1,2 @@
+from tensorflow.keras import layers
+sp = layers.builder.config("config3", "1").config("config1", "5").config("config2", "5").getOrCreate()

test-resources/python/find_import_extend_method_chain/input/method_chain.py

@@ -0,0 +1,2 @@
+from tensorflow.keras import layers
+sp = layers.builder.config("config1", "5").config("config2", "5").getOrCreate()

tests/tests.py

@@ -10,30 +10,31 @@
 # limitations under the License.
 from pathlib import Path
 from polyglot_piranha import Filter, execute_piranha, PiranhaArguments, PiranhaOutputSummary, Rule, RuleGraph, OutgoingEdges
-from os.path import join, basename
+from os.path import join, basename, dirname
 from os import listdir
 import re
 import pytest
 
 def test_piranha_rewrite():
     args = PiranhaArguments(
-        path_to_configurations="test-resources/java/feature_flag_system_1/treated/configurations",
+        path_to_configurations=join(dirname(__file__), "../test-resources/java/feature_flag_system_1/treated/configurations"),
+        # path_to_configurations="test-resources/java/feature_flag_system_1/treated/configurations",
         language="java",
         substitutions={
             "stale_flag_name": "STALE_FLAG",
             "treated": "true",
             "treated_complement": "false",
         },
-        paths_to_codebase=["test-resources/java/feature_flag_system_1/treated/input"],
+        paths_to_codebase=[join(dirname(__file__),"../test-resources/java/feature_flag_system_1/treated/input")],
         dry_run=True,
     )
 
     output_summaries = execute_piranha(args)
 
     assert len(output_summaries) == 2
     expected_paths = [
-        "test-resources/java/feature_flag_system_1/treated/input/XPFlagCleanerPositiveCases.java",
-        "test-resources/java/feature_flag_system_1/treated/input/TestEnum.java",
+        join(dirname(__file__), "../test-resources/java/feature_flag_system_1/treated/input/XPFlagCleanerPositiveCases.java"),
+        join(dirname(__file__), "../test-resources/java/feature_flag_system_1/treated/input/TestEnum.java"),
     ]
     assert all([o.path in expected_paths for o in output_summaries])
     summary: PiranhaOutputSummary
@@ -45,14 +46,14 @@ def test_piranha_rewrite():
             assert rewrite.p_match.matched_string and rewrite.p_match.matches
 
     assert is_as_expected(
-        "test-resources/java/feature_flag_system_1/treated", output_summaries
+        join(dirname(__file__), "../test-resources/java/feature_flag_system_1/treated"), output_summaries
     )
 
 def test_piranha_match_only():
     args = PiranhaArguments(
-        path_to_configurations="test-resources/java/structural_find_with_include_exclude/configurations",
+        path_to_configurations=join(dirname(__file__), "../test-resources/java/structural_find_with_include_exclude/configurations"),
         language="java",
-        paths_to_codebase=["test-resources/java/structural_find_with_include_exclude/input"],
+        paths_to_codebase=[join(dirname(__file__), "../test-resources/java/structural_find_with_include_exclude/input")],
         dry_run=True,
         exclude=["*/folder_2_1/**/*"]
     )
@@ -122,7 +123,7 @@ def test_insert_field_add_import():
         )
 
     args = PiranhaArguments(
-        paths_to_codebase=["test-resources/java/insert_field_and_import/input"],
+        paths_to_codebase=[join(dirname(__file__), "../test-resources/java/insert_field_and_import/input")],
         language="java",
         rule_graph = rule_graph,
         dry_run=True,
@@ -164,7 +165,7 @@ def test_delete_unused_field():
     )
 
     args = PiranhaArguments(
-        paths_to_codebase=["test-resources/java/delete_unused_field/input"],
+        paths_to_codebase=[join(dirname(__file__), "../test-resources/java/delete_unused_field/input")],
         language="java",
         rule_graph = rule_graph,
         dry_run=True,
@@ -173,7 +174,7 @@ def test_delete_unused_field():
     output_summaries = execute_piranha(args)
     print(output_summaries[0].content)
     assert is_as_expected(
-        "test-resources/java/delete_unused_field/", output_summaries
+        join(dirname(__file__), "../test-resources/java/delete_unused_field/"), output_summaries
     )
 
 
@@ -259,3 +260,25 @@ def _java_toplevel_mdecl_matches_anything(code_snippet: str) -> bool:
             return len(summaries) > 0
         except:
             assert False, f"Java method_declaration as top level node should not raise an error:\n{code_snippet}"
+
+def test_python_method_chain_extension():
+    args = PiranhaArguments(
+        path_to_configurations=join(dirname(__file__), "../test-resources/python/find_import_extend_method_chain/configurations"),
+        # path_to_configurations="test-resources/java/feature_flag_system_1/treated/configurations",
+        language="python",
+        substitutions={
+            "imported_name_to_check": "layers",
+            "module_name_to_check": "tensorflow",
+            "submodule_name_to_check": "keras",
+            "class_attribute": "builder",
+            "method_call_to_check": "config",
+            "method_call_to_add": "config(\"config3\", \"1\")"
+        },
+        paths_to_codebase=[join(dirname(__file__),"../test-resources/python/find_import_extend_method_chain/input")],
+        dry_run=True,
+    )
+    output_summaries = execute_piranha(args)
+    assert is_as_expected(
+        join(dirname(__file__), "../test-resources/python/find_import_extend_method_chain/"), output_summaries
+    )
+