Repository: ast-grep/ast-grep (language: rust)

Issue:
Is there any way to configure custom language support when using the Python API?
I'm not seeing anywhere in the API where an `sgconfig.yaml` can be defined. 😕 I've been exploring both the docs and the source code and am not seeing how to do it. Any suggestions?

docs I've looked at:
- https://ast-grep.github.io/guide/api-usage/py-api.html
- https://ast-grep.github.io/reference/api.html
- https://github.com/ast-grep/ast-grep/blob/0.22.3/crates/pyo3/src/py_node.rs


Discussion / hints:
It is not supported yet. Custom language will be supported after fix API.

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index f1fd000a84..6549ab9384 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -253,9 +253,11 @@ dependencies = [
  "anyhow",
  "ast-grep-config",
  "ast-grep-core",
+ "ast-grep-dynamic",
  "ast-grep-language",
  "pyo3",
  "pythonize",
+ "serde",
 ]
 
 [[package]]
diff --git a/crates/pyo3/Cargo.toml b/crates/pyo3/Cargo.toml
index 7547b1868b..d1919e3949 100644
--- a/crates/pyo3/Cargo.toml
+++ b/crates/pyo3/Cargo.toml
@@ -19,9 +19,11 @@ crate-type = ["cdylib"]
 ast-grep-core.workspace = true
 ast-grep-config.workspace = true
 ast-grep-language.workspace = true
+ast-grep-dynamic.workspace = true
 anyhow.workspace = true
 pyo3 = { version = "0.21.2", optional = true, features = ["anyhow"] }
 pythonize = { version = "0.21.1", optional = true }
+serde.workspace = true
 
 # uncomment default features when developing pyo3
 [features]
diff --git a/crates/pyo3/ast_grep_py/__init__.py b/crates/pyo3/ast_grep_py/__init__.py
index 95278b91e3..26b7ce66d2 100644
--- a/crates/pyo3/ast_grep_py/__init__.py
+++ b/crates/pyo3/ast_grep_py/__init__.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from typing import List, TypedDict,  Literal, Dict, Union, Mapping
-from .ast_grep_py import SgNode, SgRoot, Pos, Range, Edit
+from typing import List, TypedDict,  Literal, Dict, Union, Mapping, Optional
+from .ast_grep_py import SgNode, SgRoot, Pos, Range, Edit, register_dynamic_language
 
 Strictness = Union[Literal["cst"], Literal["smart"], Literal["ast"], Literal["relaxed"], Literal["signature"]]
 
@@ -55,6 +55,12 @@ class Config(TypedDict, total=False):
     utils: Dict[str, Rule]
     transform: Dict[str, Mapping]
 
+class CustomLang(TypedDict, total=False):
+  library_path: str
+  language_symbol: Optional[str]
+  meta_var_char: Optional[str]
+  expando_char: Optional[str]
+
 __all__ = [
     "Rule",
     "Config",
@@ -66,4 +72,5 @@ class Config(TypedDict, total=False):
     "Pos",
     "Range",
     "Edit",
+    "register_dynamic_language",
 ]
\ No newline at end of file
diff --git a/crates/pyo3/ast_grep_py/ast_grep_py.pyi b/crates/pyo3/ast_grep_py/ast_grep_py.pyi
index db3e73ba30..850009e52b 100644
--- a/crates/pyo3/ast_grep_py/ast_grep_py.pyi
+++ b/crates/pyo3/ast_grep_py/ast_grep_py.pyi
@@ -1,6 +1,6 @@
-from typing import List, Optional, overload, Unpack
+from typing import List, Optional, overload, Unpack, Dict
 
-from . import Rule, Config
+from . import Rule, Config, CustomLang
 
 class Pos:
     line: int
@@ -65,4 +65,6 @@ class SgNode:
 
     # Edit
     def replace(self, new_text: str) -> Edit: ...
-    def commit_edits(self, edits: List[Edit]) -> str: ...
\ No newline at end of file
+    def commit_edits(self, edits: List[Edit]) -> str: ...
+
+def register_dynamic_language(langs: Dict[str, CustomLang]): ...
\ No newline at end of file
diff --git a/crates/pyo3/src/lib.rs b/crates/pyo3/src/lib.rs
index 3b9ca59c99..0853cacb17 100644
--- a/crates/pyo3/src/lib.rs
+++ b/crates/pyo3/src/lib.rs
@@ -1,13 +1,15 @@
 #![cfg(not(test))]
 #![cfg(feature = "python")]
+mod py_lang;
 mod py_node;
 mod range;
 mod unicode_position;
+use py_lang::register_dynamic_language;
 use py_node::{Edit, SgNode};
 use range::{Pos, Range};
 
 use ast_grep_core::{AstGrep, Language, NodeMatch, StrDoc};
-use ast_grep_language::SupportLang;
+use py_lang::PyLang;
 use pyo3::prelude::*;
 
 use unicode_position::UnicodePosition;
@@ -20,12 +22,13 @@ fn ast_grep_py(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
   m.add_class::<Range>()?;
   m.add_class::<Pos>()?;
   m.add_class::<Edit>()?;
+  m.add_function(wrap_pyfunction!(register_dynamic_language, m)?)?;
   Ok(())
 }
 
 #[pyclass]
 struct SgRoot {
-  inner: AstGrep<StrDoc<SupportLang>>,
+  inner: AstGrep<StrDoc<PyLang>>,
   filename: String,
   pub(crate) position: UnicodePosition,
 }
@@ -35,7 +38,7 @@ impl SgRoot {
   #[new]
   fn new(src: &str, lang: &str) -> Self {
     let position = UnicodePosition::new(src);
-    let lang: SupportLang = lang.parse().unwrap();
+    let lang: PyLang = lang.parse().unwrap();
     let inner = lang.ast_grep(src);
     Self {
       inner,
diff --git a/crates/pyo3/src/py_lang.rs b/crates/pyo3/src/py_lang.rs
new file mode 100644
index 0000000000..7678b8eca3
--- /dev/null
+++ b/crates/pyo3/src/py_lang.rs
@@ -0,0 +1,125 @@
+use ast_grep_core::language::TSLanguage;
+use ast_grep_dynamic::{DynamicLang, Registration};
+use ast_grep_language::{Language, SupportLang};
+use serde::{Deserialize, Serialize};
+
+use pyo3::prelude::*;
+use pyo3::types::PyDict;
+use pythonize::depythonize_bound;
+
+use std::borrow::Cow;
+use std::collections::HashMap;
+use std::fmt::{Debug, Display, Formatter};
+use std::path::{Path, PathBuf};
+use std::str::FromStr;
+
+#[derive(Serialize, Deserialize, Clone)]
+pub struct CustomLang {
+  library_path: PathBuf,
+  /// the dylib symbol to load ts-language, default is `tree_sitter_{name}`
+  language_symbol: Option<String>,
+  meta_var_char: Option<char>,
+  expando_char: Option<char>,
+  // extensions: Vec<String>,
+}
+
+#[pyfunction]
+pub fn register_dynamic_language(dict: Bound<PyDict>) -> PyResult<()> {
+  let langs = depythonize_bound(dict.into_any())?;
+  let base = std::env::current_dir()?;
+  register(base, langs);
+  Ok(())
+}
+
+fn register(base: PathBuf, langs: HashMap<String, CustomLang>) {
+  let registrations = langs
+    .into_iter()
+    .map(|(name, custom)| to_registration(name, custom, &base))
+    .collect();
+  // TODO, add error handling
+  unsafe { DynamicLang::register(registrations).expect("TODO") }
+}
+
+fn to_registration(name: String, custom_lang: CustomLang, base: &Path) -> Registration {
+  let path = base.join(custom_lang.library_path);
+  let sym = custom_lang
+    .language_symbol
+    .unwrap_or_else(|| format!("tree_sitter_{name}"));
+  Registration {
+    lang_name: name,
+    lib_path: path,
+    symbol: sym,
+    meta_var_char: custom_lang.meta_var_char,
+    expando_char: custom_lang.expando_char,
+    // extensions: custom_lang.extensions,
+    extensions: vec![],
+  }
+}
+
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
+pub enum PyLang {
+  // inlined support lang expando char
+  Builtin(SupportLang),
+  Custom(DynamicLang),
+}
+#[derive(Debug)]
+pub enum PyLangErr {
+  LanguageNotSupported(String),
+}
+
+impl Display for PyLangErr {
+  fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
+    use PyLangErr::*;
+    match self {
+      LanguageNotSupported(lang) => write!(f, "{} is not supported!", lang),
+    }
+  }
+}
+
+impl std::error::Error for PyLangErr {}
+
+impl FromStr for PyLang {
+  type Err = PyLangErr;
+  fn from_str(s: &str) -> Result<Self, Self::Err> {
+    if let Ok(b) = SupportLang::from_str(s) {
+      Ok(PyLang::Builtin(b))
+    } else if let Ok(c) = DynamicLang::from_str(s) {
+      Ok(PyLang::Custom(c))
+    } else {
+      Err(PyLangErr::LanguageNotSupported(s.into()))
+    }
+  }
+}
+
+use PyLang::*;
+impl Language for PyLang {
+  fn get_ts_language(&self) -> TSLanguage {
+    match self {
+      Builtin(b) => b.get_ts_language(),
+      Custom(c) => c.get_ts_language(),
+    }
+  }
+
+  fn pre_process_pattern<'q>(&self, query: &'q str) -> Cow<'q, str> {
+    match self {
+      Builtin(b) => b.pre_process_pattern(query),
+      Custom(c) => c.pre_process_pattern(query),
+    }
+  }
+
+  #[inline]
+  fn meta_var_char(&self) -> char {
+    match self {
+      Builtin(b) => b.meta_var_char(),
+      Custom(c) => c.meta_var_char(),
+    }
+  }
+
+  #[inline]
+  fn expando_char(&self) -> char {
+    match self {
+      Builtin(b) => b.expando_char(),
+      Custom(c) => c.expando_char(),
+    }
+  }
+}
diff --git a/crates/pyo3/src/py_node.rs b/crates/pyo3/src/py_node.rs
index 838db41a5f..fe542c8bb4 100644
--- a/crates/pyo3/src/py_node.rs
+++ b/crates/pyo3/src/py_node.rs
@@ -1,9 +1,9 @@
+use crate::py_lang::PyLang;
 use crate::range::Range;
 use crate::SgRoot;
 
 use ast_grep_config::{DeserializeEnv, RuleCore, SerializableRuleCore};
 use ast_grep_core::{NodeMatch, StrDoc};
-use ast_grep_language::SupportLang;
 
 use std::collections::hash_map::DefaultHasher;
 use std::hash::{Hash, Hasher};
@@ -16,7 +16,7 @@ use pythonize::depythonize_bound;
 
 #[pyclass(mapping)]
 pub struct SgNode {
-  pub inner: NodeMatch<'static, StrDoc<SupportLang>>,
+  pub inner: NodeMatch<'static, StrDoc<PyLang>>,
   // refcount SgRoot
   pub(crate) root: Py<SgRoot>,
 }
@@ -328,7 +328,7 @@ impl SgNode {
     &self,
     config: Option<Bound<PyDict>>,
     kwargs: Option<Bound<PyDict>>,
-  ) -> PyResult<RuleCore<SupportLang>> {
+  ) -> PyResult<RuleCore<PyLang>> {
     let lang = self.inner.lang();
     let config = if let Some(config) = config {
       config_from_dict(config)?
@@ -358,10 +358,7 @@ fn config_from_rule(dict: Bound<PyDict>) -> PyResult<SerializableRuleCore> {
   })
 }
 
-fn get_matcher_from_rule(
-  lang: &SupportLang,
-  dict: Option<Bound<PyDict>>,
-) -> PyResult<RuleCore<SupportLang>> {
+fn get_matcher_from_rule(lang: &PyLang, dict: Option<Bound<PyDict>>) -> PyResult<RuleCore<PyLang>> {
   let rule = dict.ok_or_else(|| PyErr::new::<PyValueError, _>("rule must not be empty"))?;
   let env = DeserializeEnv::new(*lang);
   let config = config_from_rule(rule)?;


Test patch:
diff --git a/crates/pyo3/tests/test_register_lang.py b/crates/pyo3/tests/test_register_lang.py
new file mode 100644
index 0000000000..f3a3cbc284
--- /dev/null
+++ b/crates/pyo3/tests/test_register_lang.py
@@ -0,0 +1,35 @@
+from ast_grep_py import SgRoot, register_dynamic_language
+import platform
+
+source = """
+{
+  "test": 123
+}
+""".strip()
+
+# at the moment only test darwin arm64
+def is_arm_mac():
+    return platform.system() == 'Darwin' and platform.processor() == 'arm'
+
+def register_lang():
+    if not is_arm_mac():
+        return
+    register_dynamic_language({
+        "myjson": {
+            "library_path": "../../benches/fixtures/json-mac.so",
+            "language_symbol": "tree_sitter_json",
+        }
+    })
+
+register_lang()
+
+def test_load_json_lang():
+    if not is_arm_mac():
+        return
+    sg = SgRoot(source, "myjson")
+    root = sg.root()
+    node = root.find(pattern="123")
+    assert node
+    assert node.kind() == "number"
+    node = root.find(pattern="456")
+    assert not node
\ No newline at end of file