From b994ed1658393222b3039fb67258c0a4a4ea1032 Mon Sep 17 00:00:00 2001
From: williamfzc <178894043@qq.com>
Date: Sun, 21 Apr 2024 22:52:32 +0800
Subject: [PATCH] feat: python support

---
 .github/workflows/cargo-test.yml |  6 +++++
 Cargo.lock                       | 11 +++++++++
 Cargo.toml                       |  1 +
 src/extractor.rs                 | 41 ++++++++++++++++++++++++++++++++
 src/graph.rs                     |  9 +++++++
 src/rule.rs                      | 10 ++++++++
 6 files changed, 78 insertions(+)

diff --git a/.github/workflows/cargo-test.yml b/.github/workflows/cargo-test.yml
index 89e8834..3f1ee83 100644
--- a/.github/workflows/cargo-test.yml
+++ b/.github/workflows/cargo-test.yml
@@ -37,6 +37,8 @@ jobs:
           git clone https://github.com/typescript-eslint/typescript-eslint --depth 128
           # golang
           git clone https://github.com/gin-gonic/gin --depth 128
+          # python
+          git clone https://github.com/tiangolo/fastapi --depth 128
           
           # execute bin
           cp ./target/debug/gossiphs ./gossiphs
@@ -52,3 +54,7 @@ jobs:
           cd gin
           time ../gossiphs interactive --dry
           cd ..
+          
+          cd fastapi
+          time ../gossiphs interactive --dry
+          cd ..
diff --git a/Cargo.lock b/Cargo.lock
index 0a92881..a251f85 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -436,6 +436,7 @@ dependencies = [
  "tracing-subscriber",
  "tree-sitter",
  "tree-sitter-go",
+ "tree-sitter-python",
  "tree-sitter-rust",
  "tree-sitter-typescript",
 ]
@@ -1265,6 +1266,16 @@ dependencies = [
  "tree-sitter",
 ]
 
+[[package]]
+name = "tree-sitter-python"
+version = "0.20.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c93b1b1fbd0d399db3445f51fd3058e43d0b4dcff62ddbdb46e66550978aa5"
+dependencies = [
+ "cc",
+ "tree-sitter",
+]
+
 [[package]]
 name = "tree-sitter-rust"
 version = "0.20.4"
diff --git a/Cargo.toml b/Cargo.toml
index d0a41d3..e661cc1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,6 +19,7 @@ tree-sitter-rust = "0.20.4"
 petgraph = "0.6.4"
 tree-sitter-typescript = "0.20.5"
 tree-sitter-go = "0.20.0"
+tree-sitter-python = "0.20.4"
 serde = { version = "1.0.197", features = ["derive"] }
 indicatif = "0.17.8"
 inquire = "0.7.4"
diff --git a/src/extractor.rs b/src/extractor.rs
index 738b2b1..0170681 100644
--- a/src/extractor.rs
+++ b/src/extractor.rs
@@ -7,6 +7,7 @@ pub enum Extractor {
     Rust,
     TypeScript,
     Go,
+    Python,
 }
 
 impl Extractor {
@@ -24,6 +25,10 @@ impl Extractor {
                 let lang = &tree_sitter_go::language();
                 self._extract(f, s, lang)
             }
+            Extractor::Python => {
+                let lang = &tree_sitter_python::language();
+                self._extract(f, s, lang)
+            }
         };
     }
 
@@ -227,4 +232,40 @@ func injectV1Group(v1group *gin.RouterGroup) {
             info!("symbol: {:?} {:?}", each.name, each.kind);
         })
     }
+
+    #[test]
+    fn extract_python() {
+        let symbols = Extractor::Python.extract(
+            &String::from("abc"),
+            &String::from(
+                r#"
+def normal_fff(self, env_config: EnvConfig):
+    pass
+
+class BaseStep(object):
+    def apply(self, env_config: EnvConfig, result: ResultContext):
+        raise NotImplementedError
+
+    def name(self) -> str:
+        raise NotImplementedError
+
+    def config_name(self) -> str:
+        return self.name().replace("-", "_")
+
+    def get_mod_config(self, env_config: EnvConfig):
+        return getattr(
+            env_config._repo_config.modules,
+            self.config_name(),
+        )
+
+    def enabled(self, env_config: EnvConfig) -> bool:
+        mod_config = self.get_mod_config(env_config)
+        return mod_config.enabled
+            "#,
+            ),
+        );
+        symbols.iter().for_each(|each| {
+            info!("symbol: {:?}", each);
+        })
+    }
 }
diff --git a/src/graph.rs b/src/graph.rs
index f73cb61..eab417f 100644
--- a/src/graph.rs
+++ b/src/graph.rs
@@ -84,6 +84,15 @@ impl Graph {
                     };
                     file_contexts.push(file_context);
                 }
+                "py" => {
+                    let symbols = Extractor::Python.extract(each_file, file_content);
+                    let file_context = FileContext {
+                        // use the relative path as key
+                        path: each_file.clone(),
+                        symbols,
+                    };
+                    file_contexts.push(file_context);
+                }
                 _ => {}
             }
         }
diff --git a/src/rule.rs b/src/rule.rs
index b51301c..83f48a0 100644
--- a/src/rule.rs
+++ b/src/rule.rs
@@ -41,6 +41,16 @@ pub fn get_rule(extractor_type: &Extractor) -> Rule {
             export_grammar: r#"
 (function_declaration name: (identifier) @exported_symbol)
 (method_declaration name: (field_identifier) @exported_symbol)
+"#,
+        },
+
+        Extractor::Python => Rule {
+            import_grammar: r#"
+(identifier) @variable_name
+"#,
+            export_grammar: r#"
+(function_definition name: (identifier) @exported_symbol)
+(class_definition name: (identifier) @exported_symbol)
 "#,
         },
     }