1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

#![cfg(not(test))]
#![cfg(feature = "python")]
mod py_node;
mod range;
mod unicode_position;
use py_node::{Edit, SgNode};
use range::{Pos, Range};

use ast_grep_core::{AstGrep, Language, NodeMatch, StrDoc};
use ast_grep_language::SupportLang;
use pyo3::prelude::*;

use unicode_position::UnicodePosition;

/// A Python module implemented in Rust.
#[pymodule]
fn ast_grep_py(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
  m.add_class::<SgRoot>()?;
  m.add_class::<SgNode>()?;
  m.add_class::<Range>()?;
  m.add_class::<Pos>()?;
  m.add_class::<Edit>()?;
  Ok(())
}

#[pyclass]
struct SgRoot {
  inner: AstGrep<StrDoc<SupportLang>>,
  filename: String,
  pub(crate) position: UnicodePosition,
}

#[pymethods]
impl SgRoot {
  #[new]
  fn new(src: &str, lang: &str) -> Self {
    let position = UnicodePosition::new(src);
    let lang: SupportLang = lang.parse().unwrap();
    let inner = lang.ast_grep(src);
    Self {
      inner,
      filename: "anonymous".into(),
      position,
    }
  }

  fn root(slf: PyRef<Self>) -> SgNode {
    let tree = unsafe { &*(&slf.inner as *const AstGrep<_>) } as &'static AstGrep<_>;
    let inner = NodeMatch::from(tree.root());
    SgNode {
      inner,
      root: slf.into(),
    }
  }

  fn filename(&self) -> &str {
    &self.filename
  }
}