Merge pull request #2295 from CosmWasm/tkulik/python_codegen

Python codegen

packages/cw-schema-codegen/src/lib.rs

@@ -1,3 +1,4 @@
 pub mod go;
+pub mod python;
 pub mod rust;
 pub mod typescript;

packages/cw-schema-codegen/src/main.rs

@@ -74,7 +74,8 @@ where
             Language::Typescript => {
                 cw_schema_codegen::typescript::process_node(output, schema, node, add_imports)
             }
-            Language::Go | Language::Python => todo!(),
+            Language::Python => cw_schema_codegen::python::process_node(output, schema, node),
+            Language::Go => todo!(),
         }
     })?;
 

packages/cw-schema-codegen/src/python/mod.rs

@@ -0,0 +1,139 @@
+use self::template::{
+    EnumTemplate, EnumVariantTemplate, FieldTemplate, StructTemplate, TypeTemplate,
+};
+use std::{borrow::Cow, io};
+
+pub mod template;
+
+fn expand_node_name<'a>(
+    schema: &'a cw_schema::SchemaV1,
+    node: &'a cw_schema::Node,
+) -> Cow<'a, str> {
+    match node.value {
+        cw_schema::NodeType::Array { items } => {
+            let items = &schema.definitions[items];
+            format!("{}[]", expand_node_name(schema, items)).into()
+        }
+        cw_schema::NodeType::Float => "float".into(),
+        cw_schema::NodeType::Double => "float".into(),
+        cw_schema::NodeType::Boolean => "bool".into(),
+        cw_schema::NodeType::String => "str".into(),
+        cw_schema::NodeType::Integer { .. } => "int".into(),
+        cw_schema::NodeType::Binary => "bytes".into(),
+        cw_schema::NodeType::Optional { inner } => {
+            let inner = &schema.definitions[inner];
+            format!("typing.Optional[{}]", expand_node_name(schema, inner)).into()
+        }
+        cw_schema::NodeType::Struct(..) => node.name.as_ref().into(),
+        cw_schema::NodeType::Tuple { ref items } => {
+            let items = items
+                .iter()
+                .map(|item| expand_node_name(schema, &schema.definitions[*item]))
+                .collect::<Vec<_>>()
+                .join(", ");
+
+            format!("[{}]", items).into()
+        }
+        cw_schema::NodeType::Enum { .. } => node.name.as_ref().into(),
+
+        cw_schema::NodeType::Decimal { .. } => "decimal.Decimal".into(),
+        cw_schema::NodeType::Address => "str".into(),
+        cw_schema::NodeType::Checksum => todo!(),
+        cw_schema::NodeType::HexBinary => todo!(),
+        cw_schema::NodeType::Timestamp => todo!(),
+        cw_schema::NodeType::Unit => "None".into(),
+        _ => todo!(),
+    }
+}
+
+fn prepare_docs(desc: Option<&str>) -> Cow<'_, [Cow<'_, str>]> {
+    desc.map(|desc| desc.lines().map(Into::into).collect())
+        .unwrap_or(Cow::Borrowed(&[]))
+}
+
+pub fn process_node<O>(
+    output: &mut O,
+    schema: &cw_schema::SchemaV1,
+    node: &cw_schema::Node,
+) -> io::Result<()>
+where
+    O: io::Write,
+{
+    match node.value {
+        cw_schema::NodeType::Struct(ref sty) => {
+            let structt = StructTemplate {
+                name: node.name.clone(),
+                docs: prepare_docs(node.description.as_deref()),
+                ty: match sty {
+                    cw_schema::StructType::Unit => TypeTemplate::Unit,
+                    cw_schema::StructType::Named { ref properties } => TypeTemplate::Named {
+                        fields: properties
+                            .iter()
+                            .map(|(name, prop)| FieldTemplate {
+                                name: Cow::Borrowed(name),
+                                docs: prepare_docs(prop.description.as_deref()),
+                                ty: expand_node_name(schema, &schema.definitions[prop.value]),
+                            })
+                            .collect(),
+                    },
+                    cw_schema::StructType::Tuple { ref items } => TypeTemplate::Tuple(
+                        items
+                            .iter()
+                            .map(|item| expand_node_name(schema, &schema.definitions[*item]))
+                            .collect(),
+                    ),
+                    _ => todo!(),
+                },
+            };
+
+            writeln!(output, "{structt}")?;
+        }
+        cw_schema::NodeType::Enum { ref cases, .. } => {
+            let enumm = EnumTemplate {
+                name: node.name.clone(),
+                docs: prepare_docs(node.description.as_deref()),
+                variants: cases
+                    .iter()
+                    .map(|(name, case)| EnumVariantTemplate {
+                        name: name.clone(),
+                        docs: prepare_docs(case.description.as_deref()),
+                        ty: match case.value {
+                            cw_schema::EnumValue::Unit => TypeTemplate::Unit,
+                            cw_schema::EnumValue::Tuple { ref items } => {
+                                let items = items
+                                    .iter()
+                                    .map(|item| {
+                                        expand_node_name(schema, &schema.definitions[*item])
+                                    })
+                                    .collect();
+
+                                TypeTemplate::Tuple(items)
+                            }
+                            cw_schema::EnumValue::Named { ref properties, .. } => {
+                                TypeTemplate::Named {
+                                    fields: properties
+                                        .iter()
+                                        .map(|(name, prop)| FieldTemplate {
+                                            name: Cow::Borrowed(name),
+                                            docs: prepare_docs(prop.description.as_deref()),
+                                            ty: expand_node_name(
+                                                schema,
+                                                &schema.definitions[prop.value],
+                                            ),
+                                        })
+                                        .collect(),
+                                }
+                            }
+                            _ => todo!(),
+                        },
+                    })
+                    .collect(),
+            };
+
+            writeln!(output, "{enumm}")?;
+        }
+        _ => (),
+    }
+
+    Ok(())
+}

packages/cw-schema-codegen/src/python/template.rs

@@ -0,0 +1,41 @@
+use askama::Template;
+use std::borrow::Cow;
+
+#[derive(Clone)]
+pub struct EnumVariantTemplate<'a> {
+    pub name: Cow<'a, str>,
+    pub docs: Cow<'a, [Cow<'a, str>]>,
+    pub ty: TypeTemplate<'a>,
+}
+
+#[derive(Template)]
+#[template(escape = "none", path = "python/enum.tpl.py")]
+pub struct EnumTemplate<'a> {
+    pub name: Cow<'a, str>,
+    pub docs: Cow<'a, [Cow<'a, str>]>,
+    pub variants: Cow<'a, [EnumVariantTemplate<'a>]>,
+}
+
+#[derive(Clone)]
+pub struct FieldTemplate<'a> {
+    pub name: Cow<'a, str>,
+    pub docs: Cow<'a, [Cow<'a, str>]>,
+    pub ty: Cow<'a, str>,
+}
+
+#[derive(Clone)]
+pub enum TypeTemplate<'a> {
+    Unit,
+    Tuple(Cow<'a, [Cow<'a, str>]>),
+    Named {
+        fields: Cow<'a, [FieldTemplate<'a>]>,
+    },
+}
+
+#[derive(Template)]
+#[template(escape = "none", path = "python/struct.tpl.py")]
+pub struct StructTemplate<'a> {
+    pub name: Cow<'a, str>,
+    pub docs: Cow<'a, [Cow<'a, str>]>,
+    pub ty: TypeTemplate<'a>,
+}

packages/cw-schema-codegen/templates/python/enum.tpl.py

@@ -0,0 +1,41 @@
+# This code is @generated by cw-schema-codegen. Do not modify this manually.
+
+import typing
+import decimal
+from pydantic import BaseModel, RootModel
+
+class {{ name }}(RootModel):
+    """{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}"""
+
+{% for variant in variants %}
+{% match variant.ty %}
+{% when TypeTemplate::Unit %}
+    class {{ variant.name }}(RootModel):
+        """{% for doc in variant.docs %}
+        {{ doc }}
+        {% endfor %}"""
+        root: typing.Literal['{{ variant.name }}']
+{% when TypeTemplate::Tuple with (types) %}
+    class {{ variant.name }}(BaseModel):
+        """{% for doc in variant.docs %}
+        {{ doc }}
+        {% endfor %}"""
+        {{ variant.name }}: typing.Tuple[{{ types|join(", ") }}]
+{% when TypeTemplate::Named with { fields } %}
+    class {{ variant.name }}(BaseModel):
+        class __Inner(BaseModel):
+            """{% for doc in variant.docs %}
+            {{ doc }}
+            {% endfor %}"""
+            {% for field in fields %}
+            {{ field.name }}: {{ field.ty }}
+            """{% for doc in field.docs %}
+            {{ doc }}
+            {% endfor %}"""
+            {% endfor %}
+        {{ variant.name }}: __Inner
+{% endmatch %}
+{% endfor %}
+    root: typing.Union[ {% for variant in variants %} {{ variant.name }}, {% endfor %} ]

packages/cw-schema-codegen/templates/python/struct.tpl.py

@@ -0,0 +1,32 @@
+# This code is @generated by cw-schema-codegen. Do not modify this manually.
+
+import typing
+import decimal
+from pydantic import BaseModel, RootModel
+
+
+{% match ty %}
+{% when TypeTemplate::Unit %}
+class {{ name }}(RootModel):
+    '''{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}'''
+    root: None
+{% when TypeTemplate::Tuple with (types) %}
+class {{ name }}(RootModel):
+    '''{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}'''
+    root: typing.Tuple[{{ types|join(", ") }}]
+{% when TypeTemplate::Named with { fields } %}
+class {{ name }}(BaseModel):
+    '''{% for doc in docs %}
+    {{ doc }}
+    {% endfor %}'''
+    {% for field in fields %}
+    {{ field.name }}: {{ field.ty }}
+    '''{% for doc in field.docs %}
+    # {{ doc }}
+    {% endfor %}'''
+    {% endfor %}
+{% endmatch %}