zerodds_idl_python/

emitter.rs

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 ZeroDDS Contributors

//! IDL → Python codegen — emit logic.
//!
//! Generates, for an IDL specification, a Python module that uses the
//! `@idl_struct(typename=...)`/`@dataclass` convention of the
//! `zerodds-py` crate. Generated classes are usable directly with
//! `zerodds.IdlTopic(MyClass)`.

use std::collections::BTreeSet;
use std::fmt::Write as _;

use zerodds_idl::ast::types::{
    BitmaskDecl, BitsetDecl, CaseLabel, ConstExpr, ConstrTypeDecl, Declarator, Definition, EnumDef,
    ExceptDecl, FloatingType, Identifier, IntegerType, Literal, LiteralKind, Member, ModuleDef,
    PrimitiveType, ScopedName, SequenceType, Specification, StringType, StructDcl, StructDef,
    SwitchTypeSpec, TypeDecl, TypeSpec, TypedefDecl, UnaryOp, UnionDcl, UnionDef,
};

use crate::error::{IdlPythonError, Result};

/// Codegen options for Python modules.
#[derive(Debug, Clone, Default)]
pub struct PythonGenOptions {
    /// Optional header comment that lands at the top of the generated
    /// file. Newlines are emitted as separate `#` comment lines.
    pub header_comment: Option<String>,
}

/// Codegen entry point — IDL spec → complete Python source.
///
/// # Errors
///
/// `IdlPythonError::Unsupported` for IDL constructs that the Python
/// mapping does not (yet) support: `valuetype`, `interface`,
/// `fixed`, `map`, `any`, as well as union cases with non-literal
/// const expressions.
pub fn generate_python_module(spec: &Specification, opts: &PythonGenOptions) -> Result<String> {
    let mut imports = ImportSet::default();
    collect_imports(&spec.definitions, &mut imports)?;

    let mut out = String::new();
    out.push_str("# SPDX-License-Identifier: Apache-2.0\n");
    if let Some(c) = &opts.header_comment {
        for line in c.lines() {
            out.push_str("# ");
            out.push_str(line);
            out.push('\n');
        }
    }
    out.push_str("# Auto-generated by `zerodds-idl-python`. Do not edit by hand.\n");
    out.push('\n');

    imports.emit(&mut out);

    let mut scope: Vec<String> = Vec::new();
    emit_definitions(&mut out, &spec.definitions, &mut scope)?;
    Ok(out)
}

// ---------------------------------------------------------------------------
// Import-Sammlung
// ---------------------------------------------------------------------------

#[derive(Default)]
struct ImportSet {
    dataclass: bool,
    int_enum: bool,
    int_flag: bool,
    typing_list: bool,
    typing_type_alias: bool,
    zerodds_brands: BTreeSet<&'static str>,
    zerodds_idl_struct: bool,
    zerodds_idl_union: bool,
}

impl ImportSet {
    fn emit(&self, out: &mut String) {
        if self.dataclass {
            out.push_str("from dataclasses import dataclass\n");
        }
        if self.int_enum || self.int_flag {
            out.push_str("from enum import ");
            let mut parts: Vec<&str> = Vec::new();
            if self.int_enum {
                parts.push("IntEnum");
            }
            if self.int_flag {
                parts.push("IntFlag");
            }
            out.push_str(&parts.join(", "));
            out.push('\n');
        }
        if self.typing_list || self.typing_type_alias {
            out.push_str("from typing import ");
            let mut parts: Vec<&str> = Vec::new();
            if self.typing_list {
                parts.push("List");
            }
            if self.typing_type_alias {
                parts.push("TypeAlias");
            }
            out.push_str(&parts.join(", "));
            out.push('\n');
        }
        if self.zerodds_idl_struct || self.zerodds_idl_union || !self.zerodds_brands.is_empty() {
            out.push_str("from zerodds.idl import ");
            let mut parts: Vec<String> = Vec::new();
            if self.zerodds_idl_struct {
                parts.push("idl_struct".into());
            }
            if self.zerodds_idl_union {
                parts.push("idl_union".into());
            }
            for brand in &self.zerodds_brands {
                parts.push((*brand).to_string());
            }
            out.push_str(&parts.join(", "));
            out.push('\n');
        }
        out.push('\n');
    }
}

/// zerodds-lint: recursion-depth 32
fn collect_imports(defs: &[Definition], imports: &mut ImportSet) -> Result<()> {
    for d in defs {
        match d {
            Definition::Module(m) => collect_imports(&m.definitions, imports)?,
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Struct(StructDcl::Def(s)))) => {
                imports.dataclass = true;
                imports.zerodds_idl_struct = true;
                for m in &s.members {
                    collect_type_imports(&m.type_spec, imports)?;
                }
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Enum(_))) => {
                imports.int_enum = true;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Bitmask(_))) => {
                imports.int_flag = true;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Bitset(_))) => {
                // Bitset: emit as an IntN alias + bit constants —
                // needs only the underlying integer brand, no
                // decorator.
                imports.zerodds_brands.insert("Int64");
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Union(UnionDcl::Def(u)))) => {
                imports.zerodds_idl_union = true;
                imports.int_enum = true;
                collect_switch_type_imports(&u.switch_type, imports);
                for case in &u.cases {
                    collect_type_imports(&case.element.type_spec, imports)?;
                }
            }
            Definition::Type(TypeDecl::Typedef(td)) => {
                imports.typing_type_alias = true;
                collect_type_imports(&td.type_spec, imports)?;
            }
            Definition::Except(e) => {
                imports.dataclass = true;
                imports.zerodds_idl_struct = true;
                for m in &e.members {
                    collect_type_imports(&m.type_spec, imports)?;
                }
            }
            _ => {
                // Forward decls + const + interface + valuetype + RTI
                // vendor constructs are reported as Unsupported or
                // ignored in emit_definitions. For the import
                // pass we can silently skip them here.
            }
        }
    }
    Ok(())
}

/// zerodds-lint: recursion-depth 32
fn collect_type_imports(ts: &TypeSpec, imports: &mut ImportSet) -> Result<()> {
    match ts {
        TypeSpec::Primitive(p) => {
            imports.zerodds_brands.insert(brand_for_primitive(*p));
        }
        TypeSpec::String(s) => {
            imports.zerodds_brands.insert(brand_for_string(s));
        }
        TypeSpec::Sequence(seq) => {
            imports.typing_list = true;
            collect_type_imports(&seq.elem, imports)?;
        }
        TypeSpec::Scoped(_) => {
            // Local reference to another type in the same module —
            // no import needed.
        }
        TypeSpec::Fixed(_) | TypeSpec::Map(_) | TypeSpec::Any => {
            return Err(IdlPythonError::Unsupported(format!(
                "type spec not yet supported in python codegen: {ts:?}"
            )));
        }
    }
    Ok(())
}

fn collect_switch_type_imports(switch: &SwitchTypeSpec, imports: &mut ImportSet) {
    match switch {
        SwitchTypeSpec::Integer(i) => {
            imports.zerodds_brands.insert(brand_for_integer(*i));
        }
        SwitchTypeSpec::Boolean => {
            // bool is a Python builtin, no brand import needed — the
            // `_kind_from_annotation` fallback maps it to Bool.
        }
        SwitchTypeSpec::Octet => {
            imports.zerodds_brands.insert("Octet");
        }
        SwitchTypeSpec::Char => {
            imports.zerodds_brands.insert("Char");
        }
        SwitchTypeSpec::Scoped(_) => {
            // Probably an IntEnum — referenced via the defined type,
            // no extra brand import.
        }
    }
}

// ---------------------------------------------------------------------------
// Definitions-Emit
// ---------------------------------------------------------------------------

/// zerodds-lint: recursion-depth 32
fn emit_definitions(out: &mut String, defs: &[Definition], scope: &mut Vec<String>) -> Result<()> {
    for d in defs {
        match d {
            Definition::Module(m) => emit_module(out, m, scope)?,
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Struct(StructDcl::Def(s)))) => {
                emit_struct(out, s, scope)?;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Struct(StructDcl::Forward(_)))) => {
                // Forward declarations need no code in Python.
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Enum(e))) => {
                emit_enum(out, e, scope)?;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Bitmask(b))) => {
                emit_bitmask(out, b, scope)?;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Bitset(b))) => {
                emit_bitset(out, b, scope)?;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Union(UnionDcl::Def(u)))) => {
                emit_union(out, u, scope)?;
            }
            Definition::Type(TypeDecl::Constr(ConstrTypeDecl::Union(UnionDcl::Forward(_)))) => {
                // Forward declarations need no code in Python.
            }
            Definition::Type(TypeDecl::Typedef(td)) => {
                emit_typedef(out, td, scope)?;
            }
            Definition::Except(e) => {
                emit_exception(out, e, scope)?;
            }
            _ => {
                return Err(IdlPythonError::Unsupported(format!(
                    "definition variant not yet supported: {d:?}"
                )));
            }
        }
    }
    Ok(())
}

/// zerodds-lint: recursion-depth 32
fn emit_module(out: &mut String, m: &ModuleDef, scope: &mut Vec<String>) -> Result<()> {
    scope.push(m.name.text.clone());
    emit_definitions(out, &m.definitions, scope)?;
    scope.pop();
    Ok(())
}

fn emit_struct(out: &mut String, s: &StructDef, scope: &[String]) -> Result<()> {
    let typename = qualified_typename(&s.name, scope);
    let class_name = python_class_name(&s.name, scope);

    writeln!(out, "@idl_struct(typename=\"{typename}\")").ok();
    writeln!(out, "@dataclass").ok();
    if let Some(base) = &s.base {
        let base_class = python_scoped_name(base);
        writeln!(out, "class {class_name}({base_class}):").ok();
    } else {
        writeln!(out, "class {class_name}:").ok();
    }

    if s.members.is_empty() && s.base.is_none() {
        writeln!(out, "    pass").ok();
    } else if s.members.is_empty() {
        // With a base but no new fields: `pass` so the dataclass body
        // is syntactically valid.
        writeln!(out, "    pass").ok();
    } else {
        for m in &s.members {
            emit_members_of(out, m, scope)?;
        }
    }
    out.push('\n');
    Ok(())
}

fn emit_exception(out: &mut String, e: &ExceptDecl, scope: &[String]) -> Result<()> {
    let typename = qualified_typename(&e.name, scope);
    let class_name = python_class_name(&e.name, scope);

    writeln!(out, "@idl_struct(typename=\"{typename}\")").ok();
    writeln!(out, "@dataclass").ok();
    writeln!(out, "class {class_name}(Exception):").ok();

    if e.members.is_empty() {
        writeln!(out, "    pass").ok();
    } else {
        for m in &e.members {
            emit_members_of(out, m, scope)?;
        }
    }
    out.push('\n');
    Ok(())
}

fn emit_members_of(out: &mut String, m: &Member, scope: &[String]) -> Result<()> {
    let py_type = python_type_for(&m.type_spec, scope)?;
    for d in &m.declarators {
        let field_name = escape_python_keyword(d.name().text.as_str());
        match d {
            Declarator::Simple(_) => {
                writeln!(out, "    {field_name}: {py_type}").ok();
            }
            Declarator::Array(arr) => {
                let mut t = py_type.clone();
                for _ in 0..arr.sizes.len() {
                    t = format!("List[{t}]");
                }
                writeln!(out, "    {field_name}: {t}").ok();
            }
        }
    }
    Ok(())
}

fn emit_enum(out: &mut String, e: &EnumDef, scope: &[String]) -> Result<()> {
    let class_name = python_class_name(&e.name, scope);
    writeln!(out, "class {class_name}(IntEnum):").ok();
    if e.enumerators.is_empty() {
        writeln!(out, "    pass").ok();
    } else {
        for (idx, en) in e.enumerators.iter().enumerate() {
            writeln!(out, "    {} = {idx}", en.name.text).ok();
        }
    }
    out.push('\n');
    Ok(())
}

fn emit_bitmask(out: &mut String, b: &BitmaskDecl, scope: &[String]) -> Result<()> {
    let class_name = python_class_name(&b.name, scope);
    writeln!(out, "class {class_name}(IntFlag):").ok();
    if b.values.is_empty() {
        writeln!(out, "    pass").ok();
    } else {
        for (idx, v) in b.values.iter().enumerate() {
            writeln!(out, "    {} = 1 << {idx}", v.name.text).ok();
        }
    }
    out.push('\n');
    Ok(())
}

fn emit_bitset(out: &mut String, b: &BitsetDecl, scope: &[String]) -> Result<()> {
    // OMG bitset: a container integer with packed sub-fields. We
    // emit an Int64 alias plus a helper class with
    // SHIFT/MASK constants per named bitfield. The `zerodds-py`
    // runtime sees the alias as an Int64 brand and encodes accordingly;
    // the application reads/writes fields via the helper constants:
    //
    //   value = SensorFlags_Bits.encode(sensor_id=42, quality=2)
    //   sensor_id = (value >> SensorFlags_Bits.SENSOR_ID_SHIFT)
    //               & SensorFlags_Bits.SENSOR_ID_MASK
    let class_name = python_class_name(&b.name, scope);
    writeln!(out, "{class_name}: TypeAlias = Int64").ok();
    writeln!(out).ok();
    writeln!(out, "class {class_name}_Bits:").ok();
    if b.bitfields.is_empty() {
        writeln!(out, "    pass").ok();
    } else {
        let mut shift: u64 = 0;
        for bf in &b.bitfields {
            let width = eval_const_int(&bf.spec.width).ok_or_else(|| {
                IdlPythonError::Unsupported(format!(
                    "bitset width is not a literal integer: {:?}",
                    bf.spec.width
                ))
            })?;
            if width <= 0 {
                return Err(IdlPythonError::Unsupported(format!(
                    "bitset width must be > 0, got {width}"
                )));
            }
            let width = width as u64;
            let mask = if width >= 64 {
                u64::MAX
            } else {
                (1u64 << width) - 1
            };
            if let Some(name) = &bf.name {
                let up = name.text.to_uppercase();
                writeln!(out, "    {up}_SHIFT = {shift}").ok();
                writeln!(out, "    {up}_WIDTH = {width}").ok();
                writeln!(out, "    {up}_MASK  = 0x{mask:x}").ok();
            }
            shift += width;
        }
    }
    out.push('\n');
    Ok(())
}

fn emit_union(out: &mut String, u: &UnionDef, scope: &[String]) -> Result<()> {
    let typename = qualified_typename(&u.name, scope);
    let class_name = python_class_name(&u.name, scope);
    let discriminator_repr = switch_type_python_repr(&u.switch_type);

    // Collect cases: one entry per label in dict[int, tuple[str, Any]].
    // For a `default:` label we collect the case separately for the
    // `default=` parameter of the idl_union constructor.
    let mut case_entries: Vec<(i64, String, String)> = Vec::new();
    let mut default_entry: Option<(String, String)> = None;

    for case in &u.cases {
        let field_name = escape_python_keyword(case.element.declarator.name().text.as_str());
        let py_type = python_type_for(&case.element.type_spec, scope)?;
        for label in &case.labels {
            match label {
                CaseLabel::Value(expr) => {
                    let v = eval_const_int(expr).ok_or_else(|| {
                        IdlPythonError::Unsupported(format!(
                            "union case label is not a literal integer (scoped/enum-ref \
                             references will be supported in a follow-up): {expr:?}"
                        ))
                    })?;
                    case_entries.push((v, field_name.clone(), py_type.clone()));
                }
                CaseLabel::Default => {
                    default_entry = Some((field_name.clone(), py_type.clone()));
                }
            }
        }
    }

    writeln!(out, "{class_name} = idl_union(").ok();
    writeln!(out, "    typename=\"{typename}\",").ok();
    writeln!(out, "    discriminator={discriminator_repr},").ok();
    writeln!(out, "    cases={{").ok();
    for (label_value, field, py_type) in &case_entries {
        writeln!(out, "        {label_value}: (\"{field}\", {py_type}),").ok();
    }
    writeln!(out, "    }},").ok();
    if let Some((field, py_type)) = &default_entry {
        writeln!(out, "    default=(\"{field}\", {py_type}),").ok();
    } else {
        writeln!(out, "    default=None,").ok();
    }
    writeln!(out, ")").ok();
    out.push('\n');
    Ok(())
}

fn emit_typedef(out: &mut String, td: &TypedefDecl, scope: &[String]) -> Result<()> {
    let py_type = python_type_for(&td.type_spec, scope)?;
    for d in &td.declarators {
        let alias_name = python_class_name(d.name(), scope);
        match d {
            Declarator::Simple(_) => {
                writeln!(out, "{alias_name}: TypeAlias = {py_type}").ok();
            }
            Declarator::Array(arr) => {
                let mut t = py_type.clone();
                for _ in 0..arr.sizes.len() {
                    t = format!("List[{t}]");
                }
                writeln!(out, "{alias_name}: TypeAlias = {t}").ok();
            }
        }
    }
    out.push('\n');
    Ok(())
}

// ---------------------------------------------------------------------------
// Type-Mapping
// ---------------------------------------------------------------------------

/// zerodds-lint: recursion-depth 32
fn python_type_for(ts: &TypeSpec, scope: &[String]) -> Result<String> {
    match ts {
        TypeSpec::Primitive(p) => Ok(brand_for_primitive(*p).to_string()),
        TypeSpec::String(s) => Ok(brand_for_string(s).to_string()),
        TypeSpec::Sequence(seq) => {
            let inner = python_type_for_seq_elem(seq, scope)?;
            Ok(format!("List[{inner}]"))
        }
        TypeSpec::Scoped(name) => Ok(python_scoped_name(name)),
        TypeSpec::Fixed(_) | TypeSpec::Map(_) | TypeSpec::Any => Err(IdlPythonError::Unsupported(
            format!("type spec not yet supported in python codegen: {ts:?}"),
        )),
    }
}

/// zerodds-lint: recursion-depth 32
fn python_type_for_seq_elem(seq: &SequenceType, scope: &[String]) -> Result<String> {
    python_type_for(&seq.elem, scope)
}

fn brand_for_primitive(p: PrimitiveType) -> &'static str {
    match p {
        PrimitiveType::Boolean => "bool",
        PrimitiveType::Octet => "Octet",
        PrimitiveType::Char => "Char",
        PrimitiveType::WideChar => "WChar",
        PrimitiveType::Integer(i) => brand_for_integer(i),
        PrimitiveType::Floating(f) => match f {
            FloatingType::Float => "Float32",
            FloatingType::Double => "Float64",
            FloatingType::LongDouble => "LongDouble",
        },
    }
}

fn brand_for_integer(i: IntegerType) -> &'static str {
    match i {
        IntegerType::Short | IntegerType::Int16 => "Int16",
        IntegerType::UShort | IntegerType::UInt16 => "UInt16",
        IntegerType::Long | IntegerType::Int32 => "Int32",
        IntegerType::ULong | IntegerType::UInt32 => "UInt32",
        IntegerType::LongLong | IntegerType::Int64 => "Int64",
        IntegerType::ULongLong | IntegerType::UInt64 => "UInt64",
        IntegerType::Int8 => "Int8",
        IntegerType::UInt8 => "UInt8",
    }
}

fn brand_for_string(s: &StringType) -> &'static str {
    if s.wide { "WString" } else { "String" }
}

fn switch_type_python_repr(switch: &SwitchTypeSpec) -> String {
    match switch {
        SwitchTypeSpec::Integer(i) => brand_for_integer(*i).to_string(),
        SwitchTypeSpec::Boolean => "bool".to_string(),
        SwitchTypeSpec::Octet => "Octet".to_string(),
        SwitchTypeSpec::Char => "Char".to_string(),
        SwitchTypeSpec::Scoped(s) => python_scoped_name(s),
    }
}

fn python_scoped_name(name: &ScopedName) -> String {
    name.parts
        .iter()
        .map(|p| p.text.clone())
        .collect::<Vec<_>>()
        .join("_")
}

fn qualified_typename(name: &Identifier, scope: &[String]) -> String {
    if scope.is_empty() {
        name.text.clone()
    } else {
        format!("{}::{}", scope.join("::"), name.text)
    }
}

fn python_class_name(name: &Identifier, scope: &[String]) -> String {
    if scope.is_empty() {
        name.text.clone()
    } else {
        let mut parts = scope.to_vec();
        parts.push(name.text.clone());
        parts.join("_")
    }
}

fn escape_python_keyword(name: &str) -> String {
    const PYTHON_KEYWORDS: &[&str] = &[
        "False", "None", "True", "and", "as", "assert", "async", "await", "break", "class",
        "continue", "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
        "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass", "raise", "return",
        "try", "while", "with", "yield", "match", "case",
    ];
    if PYTHON_KEYWORDS.contains(&name) {
        format!("{name}_")
    } else {
        name.to_string()
    }
}

// ---------------------------------------------------------------------------
// Const-expression evaluation (limited: integer literals + unary +/- only)
// ---------------------------------------------------------------------------

/// zerodds-lint: recursion-depth 32
fn eval_const_int(expr: &ConstExpr) -> Option<i64> {
    match expr {
        ConstExpr::Literal(Literal {
            kind: LiteralKind::Integer,
            raw,
            ..
        }) => parse_integer_literal(raw),
        ConstExpr::Unary { op, operand, .. } => {
            let inner = eval_const_int(operand)?;
            match op {
                UnaryOp::Plus => Some(inner),
                UnaryOp::Minus => Some(-inner),
                UnaryOp::BitNot => Some(!inner),
            }
        }
        _ => None,
    }
}

fn parse_integer_literal(raw: &str) -> Option<i64> {
    let s = raw.trim();
    // OMG IDL allows hex (0x...), octal (0...) and decimal literals.
    // Suffixes like `L`/`u` we ignore generously — the parser
    // has already validated that this is an integer literal.
    let s = s.trim_end_matches(['L', 'l', 'u', 'U']);
    if let Some(hex) = s.strip_prefix("0x").or_else(|| s.strip_prefix("0X")) {
        i64::from_str_radix(hex, 16).ok()
    } else if let Some(rest) = s.strip_prefix("-0x").or_else(|| s.strip_prefix("-0X")) {
        i64::from_str_radix(rest, 16).ok().map(|n| -n)
    } else if s.starts_with('0') && s.len() > 1 && !s.contains(|c: char| !c.is_ascii_digit()) {
        // Octal — intentionally strict, because OMG IDL accepts octal.
        i64::from_str_radix(&s[1..], 8).ok()
    } else {
        s.parse::<i64>().ok()
    }
}