mod types;

use std::fs::OpenOptions;

use std::io::{Read, Write};

const STUB_TOP: &str = r#"

let stub = format!("\n{} = dict[str, Any]\n", name);

let mut file = OpenOptions::new()

.create(true)

.write(true)

.append(true)

.read(true)

.open("python/pgml/pgml.pyi")

.unwrap();

let mut contents = String::new();

file.read_to_string(&mut contents)

.expect("Unable to read stubs file for python");

if !contents.contains(&stub) {

file.write_all(stub.as_bytes())

.expect("Unable to write stubs file for python");

}

let python_class_name = wrapped_type_ident.to_string();

let mut stubs = format!("\nclass {}:\n", python_class_name);

let signature = quote! {

pub fn #method_ident<'a>(#(#method_arguments),*) -> #output_type

};

let p1 = if method.is_async { "async def" } else { "def" };

let p2 = match method_ident.to_string().as_str() {

"new" => "__init__".to_string(),

_ => method_ident.to_string(),

};

let p3 = method

.method_arguments

.iter()

.map(|a| format!("{}: {}", a.0, get_python_type(&a.1)))

.collect::<Vec<String>>()

.join(", ");

let p4 = match &method.output_type {

OutputType::Result(v) | OutputType::Other(v) => get_python_type(v),

OutputType::Default => "None".to_string(),

};

stubs.push_str(&format!("\t{} {}(self, {}) -> {}", p1, p2, p3, p4));

stubs.push_str("\n\t\t...\n");

#signature

let middle = if method.is_async {

quote! {

get_or_set_runtime().block_on(#wrapped_type_ident::new(#(#wrapper_arguments),*))

}

} else {

quote! {

#wrapped_type_ident::new(#(#wrapper_arguments),*)

}

};

let middle = if let OutputType::Result(_r) = method.output_type {

quote! {

let x = match #middle {

};

}

} else {

quote! {

let x = #middle;

}

};

let middle = quote! {

#middle

Ok(#name_ident::from(x))

};

let x = #middle;

let mut file = OpenOptions::new()

.create(true)

.write(true)

.append(true)

.read(true)

.open("python/pgml/pgml.pyi")

.unwrap();

let mut contents = String::new();

file.read_to_string(&mut contents)

.expect("Unable to read stubs file for python");

if !contents.contains("A12BECOD") {

file.write_all(STUB_TOP.as_bytes())

.expect("Unable to write stubs file for python");

}

if !contents.contains(&format!("class {}:", python_class_name)) {

file.write_all(stubs.as_bytes())

.expect("Unable to write stubs file for python");

}

fn convert_output_type_convert_from_python(

Some(format_ident!("Self").into_token_stream()),

fn get_python_type(ty: &SupportedType) -> String {

match ty {

SupportedType::Reference(r) => get_python_type(r),

SupportedType::S => "Self".to_string(),

SupportedType::str | SupportedType::String => "str".to_string(),

SupportedType::Option(o) => format!(

"Optional[{}] = {}",

get_python_type(o),

get_type_for_optional(o)

),

SupportedType::Vec(v) => format!("List[{}]", get_python_type(v)),

SupportedType::HashMap((k, v)) => {

format!("Dict[{}, {}]", get_python_type(k), get_python_type(v))

}

SupportedType::Tuple(t) => {

let mut types = Vec::new();

for ty in t {

types.push(get_python_type(ty));

}

// Rust's unit type is represented as an empty tuple

if types.is_empty() {

"None".to_string()

} else {

format!("tuple[{}]", types.join(", "))

}

}

SupportedType::i64 => "int".to_string(),

SupportedType::f64 => "float".to_string(),

// Our own types

t @ SupportedType::Database

| t @ SupportedType::Collection

| t @ SupportedType::Splitter => t.to_string(),

// Add more types as required

_ => "Any".to_string(),

}

fn get_type_for_optional(ty: &SupportedType) -> String {

match ty {

SupportedType::Reference(r) => get_type_for_optional(r),

SupportedType::str | SupportedType::String => {

"\"Default set in Rust. Please check the documentation.\"".to_string()

}

SupportedType::HashMap(_) => "{}".to_string(),

SupportedType::Vec(_) => "[]".to_string(),

SupportedType::i64 => 1.to_string(),

SupportedType::f64 => 1.0.to_string(),

_ => panic!("Type not yet supported for optional python stub: {:?}", ty),

}

SupportedType::Tuple(t) => {

let mut types = Vec::new();

for ty in t {

types.push(ty.to_string());

}

format!("({})", types.join(","))

Here is a brief outline of how to use this crate and specifically add new Python classes.

There are three main macros to know about:

- `custom_derive`

- `custom_methods`

- `custom_into_py`

`custom_derive` is used when defining a new struct that you want to be available as a Python class. This macro automatically creates a wrapper for the struct postfixing the name with `Python`. For example, the following code:

Creates another struct:

You must currently implement `Debug` and `Clone` on the structs you use `custom_derive` on.

`custom_methods` is used on the impl block for a struct you want to be available as a Python class. This macro automatically creates methods that work seamlessly with pyO3. For example, the following code:

Produces similar code to the following:

Note that the macro only works on methods marked with `pub`;

`custom_into_py` is used when we want to seamlessly return Rust structs as Python dictionaries. For example, let's say we have the following code:

The `custom_into_py` macro automatically generates the following code for us:

Implementing `IntoPy` allows pyo3 to seamlessly convert between Rust and python. Note that Python users calling `get_text_splitters` will receive a list of dictionaries.

Be aware that the only pyo3 specific code in this crate is the `pymodule` invocation in `lib.rs`. Everything else is handled by `pgml-macros`. If you want to expose a Python Class directly on the Python module you have to add it in the `pymodule` invocation. For example, if you wanted to expose `TestStruct` so Python module users could access it directly on `pgml`, you could do the following:

Now Python users can access it like so:

For local development, install [maturin](https://github.com/PyO3/maturin) and run:

You can now run the tests in `python/test.py`.