Movatterモバイル変換


[0]ホーム

URL:


Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Sign up
Appearance settings

Generate python stub files in proc macros#730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
SilasMarvin merged 7 commits intomasterfromsilas-stub-generation
Jun 13, 2023
Merged
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletionpgml-sdks/rust/pgml-macros/src/lib.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
use syn::{parse_macro_input, DeriveInput, ItemImpl};

mod common;
mod types;
mod python;
mod types;

#[proc_macro_derive(custom_derive)]
pub fn custom_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
Expand Down
161 changes: 147 additions & 14 deletionspgml-sdks/rust/pgml-macros/src/python.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,17 @@
use quote::{format_ident, quote, ToTokens};
use std::fs::OpenOptions;
use std::io::{Read, Write};
use syn::{visit::Visit, DeriveInput, ItemImpl, Type};

use crate::common::{AttributeArgs, GetImplMethod};
use crate::types::{GetSupportedType, OutputType, SupportedType};

const STUB_TOP: &str = r#"
# Top of file key: A12BECOD!
from typing import List, Dict, Optional, Self, Any

"#;

pub fn generate_into_py(parsed: DeriveInput) -> proc_macro::TokenStream {
let name = parsed.ident;
let fields_named = match parsed.data {
Expand All@@ -29,6 +37,23 @@ pub fn generate_into_py(parsed: DeriveInput) -> proc_macro::TokenStream {
}
}).collect();

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 expanded = quote! {
impl IntoPy<PyObject> for #name {
fn into_py(self, py: Python<'_>) -> PyObject {
Expand DownExpand Up@@ -76,6 +101,9 @@ pub fn generate_python_methods(
};
let name_ident = format_ident!("{}Python", wrapped_type_ident);

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

// Iterate over the items - see: https://docs.rs/syn/latest/syn/enum.ImplItem.html
for item in parsed.items {
// We only create methods for functions listed in the attribute args
Expand DownExpand Up@@ -105,26 +133,61 @@ pub fn generate_python_methods(
OutputType::Default => (None, None),
};

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");

// The new function for pyO3 requires some unique syntax
let (signature, middle) = if method_ident == "new" {
let signature = quote! {
#[new]
pub fn new<'a>(#(#method_arguments),*) -> #output_type
#signature
};
let middle = quote! {
let runtime = get_or_set_runtime();
let x = match runtime.block_on(#wrapped_type_ident::new(#(#wrapper_arguments),*)) {
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 {
Ok(m) => m,
Err(e) => return Err(PyErr::new::<pyo3::exceptions::PyException, _>(e.to_string()))

};
Ok(#name_ident::from(x))
};
};
}
} else {
quote! {
let x = #middle;
}
};
let middle = quote! {
#middle
Ok(#name_ident::from(x))
};
(signature, middle)
} else {
let signature = quote! {
pub fn #method_ident<'a>(#(#method_arguments),*) -> #output_type
};
let middle = quote! {
#method_ident(#(#wrapper_arguments),*)
};
Expand All@@ -146,7 +209,7 @@ pub fn generate_python_methods(
}
} else {
quote! {
let x = middle;
let x =#middle;
}
};
let middle = if let Some(convert) = convert_from {
Expand DownExpand Up@@ -181,6 +244,25 @@ pub fn generate_python_methods(
});
}

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");
}

proc_macro::TokenStream::from(quote! {
#[pymethods]
impl #name_ident {
Expand DownExpand Up@@ -241,7 +323,7 @@ fn convert_method_wrapper_arguments(
}
}

pubfn convert_output_type_convert_from_python(
fn convert_output_type_convert_from_python(
ty: &SupportedType,
method: &GetImplMethod,
) -> (
Expand All@@ -251,7 +333,7 @@ pub fn convert_output_type_convert_from_python(
let (output_type, convert_from) = match ty {
SupportedType::S => (
Some(quote! {PyResult<Self>}),
Some(format_ident!("{}", method.method_ident).into_token_stream()),
Some(format_ident!("Self").into_token_stream()),
),
t @ SupportedType::Database | t @ SupportedType::Collection => (
Some(quote! {PyResult<&'a PyAny>}),
Expand All@@ -271,3 +353,54 @@ pub fn convert_output_type_convert_from_python(
(output_type, convert_from)
}
}

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),
}
}
12 changes: 6 additions & 6 deletionspgml-sdks/rust/pgml-macros/src/types.rs
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -36,12 +36,12 @@ impl ToString for SupportedType {
SupportedType::HashMap((k, v)) => {
format!("HashMap<{},{}>", k.to_string(), v.to_string())
}
SupportedType::Tuple(v) => {
let mutoutput =String::new();
v.iter().for_each(|ty| {
output.push_str(&format!("{},",ty.to_string()));
});
format!("({})",output)
SupportedType::Tuple(t) => {
let muttypes =Vec::new();
for ty in t {
types.push(ty.to_string());
}
format!("({})",types.join(","))
}
SupportedType::S => "Self".to_string(),
SupportedType::Option(v) => format!("Option<{}>", v.to_string()),
Expand Down
135 changes: 135 additions & 0 deletionspgml-sdks/rust/pgml/README.md
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -1 +1,136 @@
# Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone
# How to use this crate

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
`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:
```
#[derive(custom_derive, Debug, Clone)]
pub struct TestStruct {
pub name: String
}
```

Creates another struct:

```
pub struct TestStructPython {
pub wrapped: TestStruct
}
```

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

## custom_methods
`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:
```
#[custom_methods(new, get_name)]
impl TestStruct {
pub fn new(name: String) -> Self {
Self { name }
}
pub fn get_name(&self) -> String {
self.name.clone()
}
}
```

Produces similar code to the following:
```
impl TestStruct {
pub fn new(name: String) -> Self {
Self { name }
}
pub fn get_name(&self) -> String {
self.name.clone()
}
}

impl TestStructPython {
pub fn new<'a>(name: String, py: Python<'a>) -> PyResult<Self> {
let x = TestStruct::new(name);
Ok(TestStructPython::from(x))
}
pub fn get_name<'a>(&self, py: Python<'a>) -> PyResult<String> {
let x = self.wrapped.get_name();
Ok(x)
}
}
```

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

## custom_into_py
`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:
```
#[derive(custom_into_py, FromRow, Debug, Clone)]
pub struct Splitter {
pub id: i64,
pub created_at: DateTime<Utc>,
pub name: String,
pub parameters: Json<HashMap<String, String>>,
}

pub async fn get_text_splitters(&self) -> anyhow::Result<Vec<Splitter>> {
Ok(sqlx::query_as(&query_builder!(
"SELECT * from %s",
self.splitters_table_name
))
.fetch_all(self.pool.borrow())
.await?)
}

```

The `custom_into_py` macro automatically generates the following code for us:
```
impl IntoPy<PyObject> for Splitter {
fn into_py(self, py: Python<'_>) -> PyObject {
let dict = PyDict::new(py);
dict.set_item("id", self.id)
.expect("Error setting python value in custom_into_py proc_macro");
dict.set_item("created_at", self.created_at.timestamp())
.expect("Error setting python value in custom_into_py proc_macro");
dict.set_item("name", self.name)
.expect("Error setting python value in custom_into_py proc_macro");
dict.set_item("parameters", self.parameters.0)
.expect("Error setting python value in custom_into_py proc_macro");
dict.into()
}
}
```

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.

## Other Noteworthy Things

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:
```
#[pymodule]
fn pgml(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<TestStructPython>()?;
Ok(())
}
```

Now Python users can access it like so:
```
import pgml

t = pgml.TestStruct("test")
print(t.get_name())

```

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

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

[8]ページ先頭

©2009-2025 Movatter.jp