Skip to main content

scythe_codegen/backends/
python_common.rs

1use scythe_core::errors::{ErrorCode, ScytheError};
2
3/// Supported Python row type styles for generated code.
4#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
5pub enum PythonRowType {
6    #[default]
7    Dataclass,
8    Pydantic,
9    Msgspec,
10}
11
12impl PythonRowType {
13    /// Parse a row_type option string into a `PythonRowType`.
14    pub fn from_option(value: &str) -> Result<Self, ScytheError> {
15        match value {
16            "dataclass" => Ok(Self::Dataclass),
17            "pydantic" => Ok(Self::Pydantic),
18            "msgspec" => Ok(Self::Msgspec),
19            _ => Err(ScytheError::new(
20                ErrorCode::InternalError,
21                format!(
22                    "invalid row_type '{}': expected 'dataclass', 'pydantic', or 'msgspec'",
23                    value
24                ),
25            )),
26        }
27    }
28
29    /// Returns the import line for the row type.
30    pub fn import_line(self) -> &'static str {
31        match self {
32            Self::Dataclass => "from dataclasses import dataclass",
33            Self::Pydantic => "from pydantic import BaseModel",
34            Self::Msgspec => "import msgspec",
35        }
36    }
37
38    /// Whether the row type import is a stdlib import (vs third-party).
39    pub fn is_stdlib_import(self) -> bool {
40        matches!(self, Self::Dataclass)
41    }
42
43    /// Build a sorted third-party import block combining the row type import
44    /// with the given library import line.
45    ///
46    /// isort rules: bare `import` statements come before `from` statements,
47    /// both groups sorted by module name.
48    pub fn sorted_third_party_imports(self, library_import: &str) -> String {
49        let row_import = self.import_line();
50        let row_is_bare = row_import.starts_with("import ");
51        let lib_is_bare = library_import.starts_with("import ");
52
53        match (row_is_bare, lib_is_bare) {
54            // Both bare imports or both from imports: sort by module name.
55            (true, true) | (false, false) => {
56                if row_import < library_import {
57                    format!("{row_import}\n{library_import}")
58                } else {
59                    format!("{library_import}\n{row_import}")
60                }
61            }
62            // Row is bare, library is from: bare comes first.
63            (true, false) => format!("{row_import}\n{library_import}"),
64            // Row is from, library is bare: bare comes first.
65            (false, true) => format!("{library_import}\n{row_import}"),
66        }
67    }
68
69    /// Returns the decorator line (for dataclass) or empty string (for others).
70    pub fn decorator(self) -> &'static str {
71        match self {
72            Self::Dataclass => "@dataclass\n",
73            Self::Pydantic | Self::Msgspec => "",
74        }
75    }
76
77    /// Returns the class definition line including the class name.
78    pub fn class_def(self, class_name: &str) -> String {
79        match self {
80            Self::Dataclass => format!("class {}:", class_name),
81            Self::Pydantic => format!("class {}(BaseModel):", class_name),
82            Self::Msgspec => format!("class {}(msgspec.Struct):", class_name),
83        }
84    }
85}