Support native python enum generation (#446)

* support native python enum generation

* support native python enum generation

* Update CHANGES.rst with PR credit

Added PR credit for enum generation feature.

* Remove irrational logic

* Remove Logic

* test improvements

* Remove unneeded test

* Reinstate synthetic enum generation

* CHANGES.rst improvements

* Update CHANGES.rst

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* PR Fixes

* Fix RST formatting

* PR Fixes

* PR Fixes

* PR Fixes

* Revert PR Fix

* Update CHANGES.rst

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* PR Fixes

* Fix CHANGES.rst

* Rework CHANGES.rst

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Update src/sqlacodegen/generators.py

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* PR Fixes

* Minor cleanup

* PR Fixes

* PR Fixes

* Update CHANGES.rst

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>

* Reformatted

---------

Co-authored-by: Alex Grönholm <alex.gronholm@nextday.fi>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

CHANGES.rst

@@ -1,6 +1,17 @@
 Version history
 ===============
 
+**4.0.0rc1**
+
+- **BACKWARD INCOMPATIBLE** ``TablesGenerator.render_column_type()`` was changed to
+  receive the ``Column`` object instead of the column type object as its sole argument
+- Added Python enum generation for native database ENUM types (e.g., PostgreSQL / MySQL ENUM).
+  Retained synthetic Python enum generation from CHECK constraints with
+  IN clauses (e.g., ``column IN ('val1', 'val2', ...)``). Use ``--options nonativeenums`` to
+  disable enum generation for native database enums. Use ``--options nosyntheticenums`` to
+  disable enum generation for synthetic database enums (VARCHAR columns with check constraints).
+  (PR by @sheinbergon)
+
 **3.2.0**
 
 - Dropped support for Python 3.9

README.rst

@@ -106,6 +106,8 @@ values must be delimited by commas, e.g. ``--options noconstraints,nobidi``):
   * ``noconstraints``: ignore constraints (foreign key, unique etc.)
   * ``nocomments``: ignore table/column comments
   * ``noindexes``: ignore indexes
+  * ``nonativeenums``: don't generate Python enum classes for native database ENUM types (e.g., PostgreSQL ENUM); use plain string mapping instead
+  * ``nosyntheticenums``: don't generate Python enum classes from CHECK constraints with IN clauses (e.g., ``column IN ('value1', 'value2', ...)``); preserves CHECK constraints as-is
   * ``noidsuffix``: prevent the special naming logic for single column many-to-one
     and one-to-one relationships (see `Relationship naming logic`_ for details)
   * ``include_dialect_options``: render a table' dialect options, such as ``starrocks_partition`` for StarRocks' specific options.

src/sqlacodegen/generators.py

@@ -123,6 +123,8 @@ class TablesGenerator(CodeGenerator):
         "noindexes",
         "noconstraints",
         "nocomments",
+        "nonativeenums",
+        "nosyntheticenums",
         "include_dialect_options",
         "keep_dialect_types",
     }
@@ -148,6 +150,11 @@ def __init__(
         # Keep dialect-specific types instead of adapting to generic SQLAlchemy types
         self.keep_dialect_types: bool = "keep_dialect_types" in self.options
 
+        # Track Python enum classes: maps (table_name, column_name) -> enum_class_name
+        self.enum_classes: dict[tuple[str, str], str] = {}
+        # Track enum values: maps enum_class_name -> list of values
+        self.enum_values: dict[str, list[str]] = {}
+
     @property
     def views_supported(self) -> bool:
         return True
@@ -192,19 +199,22 @@ def generate(self) -> str:
         models: list[Model] = self.generate_models()
 
         # Render module level variables
-        variables = self.render_module_variables(models)
-        if variables:
+        if variables := self.render_module_variables(models):
             sections.append(variables + "\n")
 
+        # Render enum classes
+        if enum_classes := self.render_enum_classes():
+            sections.append(enum_classes + "\n")
+
         # Render models
-        rendered_models = self.render_models(models)
-        if rendered_models:
+        if rendered_models := self.render_models(models):
             sections.append(rendered_models)
 
         # Render collected imports
         groups = self.group_imports()
-        imports = "\n\n".join("\n".join(line for line in group) for group in groups)
-        if imports:
+        if imports := "\n\n".join(
+            "\n".join(line for line in group) for group in groups
+        ):
             sections.insert(0, imports)
 
         return "\n\n".join(sections) + "\n"
@@ -467,7 +477,7 @@ def render_column(
         # Render the column type if there are no foreign keys on it or any of them
         # points back to itself
         if not dedicated_fks or any(fk.column is column for fk in dedicated_fks):
-            args.append(self.render_column_type(column.type))
+            args.append(self.render_column_type(column))
 
         for fk in dedicated_fks:
             args.append(self.render_constraint(fk))
@@ -528,10 +538,21 @@ def render_column_callable(self, is_table: bool, *args: Any, **kwargs: Any) -> s
         else:
             return render_callable("mapped_column", *args, kwargs=kwargs)
 
-    def render_column_type(self, coltype: TypeEngine[Any]) -> str:
+    def render_column_type(self, column: Column[Any]) -> str:
+        column_type = column.type
+        # Check if this is an enum column with a Python enum class
+        if isinstance(column_type, Enum) and column is not None:
+            if enum_class_name := self.enum_classes.get(
+                (column.table.name, column.name)
+            ):
+                # Import SQLAlchemy Enum (will be handled in collect_imports)
+                self.add_import(Enum)
+                # Return the Python enum class as the type parameter
+                return f"Enum({enum_class_name})"
+
         args = []
         kwargs: dict[str, Any] = {}
-        sig = inspect.signature(coltype.__class__.__init__)
+        sig = inspect.signature(column_type.__class__.__init__)
         defaults = {param.name: param.default for param in sig.parameters.values()}
         missing = object()
         use_kwargs = False
@@ -543,7 +564,7 @@ def render_column_type(self, coltype: TypeEngine[Any]) -> str:
                 use_kwargs = True
                 continue
 
-            value = getattr(coltype, param.name, missing)
+            value = getattr(column_type, param.name, missing)
 
             if isinstance(value, (JSONB, JSON)):
                 # Remove astext_type if it's the default
@@ -577,28 +598,28 @@ def render_column_type(self, coltype: TypeEngine[Any]) -> str:
             ),
             None,
         )
-        if vararg and hasattr(coltype, vararg):
-            varargs_repr = [repr(arg) for arg in getattr(coltype, vararg)]
+        if vararg and hasattr(column_type, vararg):
+            varargs_repr = [repr(arg) for arg in getattr(column_type, vararg)]
             args.extend(varargs_repr)
 
         # These arguments cannot be autodetected from the Enum initializer
-        if isinstance(coltype, Enum):
+        if isinstance(column_type, Enum):
             for colname in "name", "schema":
-                if (value := getattr(coltype, colname)) is not None:
+                if (value := getattr(column_type, colname)) is not None:
                     kwargs[colname] = repr(value)
 
-        if isinstance(coltype, (JSONB, JSON)):
+        if isinstance(column_type, (JSONB, JSON)):
             # Remove astext_type if it's the default
             if (
-                isinstance(coltype.astext_type, Text)
-                and coltype.astext_type.length is None
+                isinstance(column_type.astext_type, Text)
+                and column_type.astext_type.length is None
             ):
                 del kwargs["astext_type"]
 
         if args or kwargs:
-            return render_callable(coltype.__class__.__name__, *args, kwargs=kwargs)
+            return render_callable(column_type.__class__.__name__, *args, kwargs=kwargs)
         else:
-            return coltype.__class__.__name__
+            return column_type.__class__.__name__
 
     def render_constraint(self, constraint: Constraint | ForeignKey) -> str:
         def add_fk_options(*opts: Any) -> None:
@@ -709,6 +730,81 @@ def find_free_name(
 
         return name
 
+    def _enum_name_to_class_name(self, enum_name: str) -> str:
+        """Convert a database enum name to a Python class name (PascalCase)."""
+        return "".join(part.capitalize() for part in enum_name.split("_") if part)
+
+    def _create_enum_class(
+        self, table_name: str, column_name: str, values: list[str]
+    ) -> str:
+        """
+        Create a Python enum class name and register it.
+
+        Returns the enum class name to use in generated code.
+        """
+        # Generate enum class name from table and column names
+        # Convert to PascalCase: user_status -> UserStatus
+        base_name = "".join(
+            part.capitalize()
+            for part in table_name.split("_") + column_name.split("_")
+            if part
+        )
+
+        # Ensure uniqueness
+        enum_class_name = base_name
+        for counter in count(1):
+            if enum_class_name not in self.enum_values:
+                break
+
+            # Check if it's the same enum (same values)
+            if self.enum_values[enum_class_name] == values:
+                # Reuse existing enum class
+                return enum_class_name
+
+            enum_class_name = f"{base_name}{counter}"
+
+        # Register the new enum class
+        self.enum_values[enum_class_name] = values
+        return enum_class_name
+
+    def render_enum_classes(self) -> str:
+        """Render Python enum class definitions."""
+        if not self.enum_values:
+            return ""
+
+        self.add_module_import("enum")
+
+        enum_defs = []
+        for enum_class_name, values in sorted(self.enum_values.items()):
+            # Create enum members with valid Python identifiers
+            members = []
+            for value in values:
+                # Unescape SQL escape sequences (e.g., \' -> ')
+                # The value from the CHECK constraint has SQL escaping
+                unescaped_value = value.replace("\\'", "'").replace("\\\\", "\\")
+
+                # Create a valid identifier from the enum value
+                member_name = _re_invalid_identifier.sub("_", unescaped_value).upper()
+                if not member_name:
+                    member_name = "EMPTY"
+                elif member_name[0].isdigit():
+                    member_name = "_" + member_name
+                elif iskeyword(member_name):
+                    member_name += "_"
+                #
+                # # Re-escape for Python string literal
+                # python_escaped = unescaped_value.replace("\\", "\\\\").replace(
+                #     "'", "\\'"
+                # )
+                members.append(f"    {member_name} = {unescaped_value!r}")
+
+            enum_def = f"class {enum_class_name}(str, enum.Enum):\n" + "\n".join(
+                members
+            )
+            enum_defs.append(enum_def)
+
+        return "\n\n\n".join(enum_defs)
+
     def fix_column_types(self, table: Table) -> None:
         """Adjust the reflected column types."""
         # Detect check constraints for boolean and enum columns
@@ -718,34 +814,74 @@ def fix_column_types(self, table: Table) -> None:
 
                 # Turn any integer-like column with a CheckConstraint like
                 # "column IN (0, 1)" into a Boolean
-                match = _re_boolean_check_constraint.match(sqltext)
-                if match:
-                    colname_match = _re_column_name.match(match.group(1))
-                    if colname_match:
+                if match := _re_boolean_check_constraint.match(sqltext):
+                    if colname_match := _re_column_name.match(match.group(1)):
                         colname = colname_match.group(3)
                         table.constraints.remove(constraint)
                         table.c[colname].type = Boolean()
                         continue
 
-                # Turn any string-type column with a CheckConstraint like
-                # "column IN (...)" into an Enum
-                match = _re_enum_check_constraint.match(sqltext)
-                if match:
-                    colname_match = _re_column_name.match(match.group(1))
-                    if colname_match:
-                        colname = colname_match.group(3)
-                        items = match.group(2)
-                        if isinstance(table.c[colname].type, String):
-                            table.constraints.remove(constraint)
-                            if not isinstance(table.c[colname].type, Enum):
-                                options = _re_enum_item.findall(items)
-                                table.c[colname].type = Enum(
-                                    *options, native_enum=False
+                # Turn VARCHAR columns with CHECK constraints like "column IN ('a', 'b')"
+                # into synthetic Enum types with Python enum classes
+                if (
+                    "nosyntheticenums" not in self.options
+                    and (match := _re_enum_check_constraint.match(sqltext))
+                    and (colname_match := _re_column_name.match(match.group(1)))
+                ):
+                    colname = colname_match.group(3)
+                    items = match.group(2)
+                    if isinstance(table.c[colname].type, String) and not isinstance(
+                        table.c[colname].type, Enum
+                    ):
+                        options = _re_enum_item.findall(items)
+                        # Create Python enum class
+                        enum_class_name = self._create_enum_class(
+                            table.name, colname, options
+                        )
+                        self.enum_classes[(table.name, colname)] = enum_class_name
+                        # Convert to Enum type but KEEP the constraint
+                        table.c[colname].type = Enum(*options, native_enum=False)
+                        continue
+
+        for column in table.c:
+            # Handle native database Enum types (e.g., PostgreSQL ENUM)
+            if (
+                "nonativeenums" not in self.options
+                and isinstance(column.type, Enum)
+                and column.type.enums
+            ):
+                if column.type.name:
+                    # Named enum - create shared enum class if not already created
+                    if (table.name, column.name) not in self.enum_classes:
+                        # Check if we've already created an enum for this name
+                        existing_class = None
+                        for (t, c), cls in self.enum_classes.items():
+                            if cls == self._enum_name_to_class_name(column.type.name):
+                                existing_class = cls
+                                break
+
+                        if existing_class:
+                            enum_class_name = existing_class
+                        else:
+                            # Create new enum class from the enum's name
+                            enum_class_name = self._enum_name_to_class_name(
+                                column.type.name
+                            )
+                            # Register the enum values if not already registered
+                            if enum_class_name not in self.enum_values:
+                                self.enum_values[enum_class_name] = list(
+                                    column.type.enums
                                 )
 
-                            continue
+                        self.enum_classes[(table.name, column.name)] = enum_class_name
+                else:
+                    # Unnamed enum - create enum class per column
+                    if (table.name, column.name) not in self.enum_classes:
+                        enum_class_name = self._create_enum_class(
+                            table.name, column.name, list(column.type.enums)
+                        )
+                        self.enum_classes[(table.name, column.name)] = enum_class_name
 
-        for column in table.c:
             if not self.keep_dialect_types:
                 try:
                     column.type = self.get_adapted_type(column.type)
@@ -1326,6 +1462,14 @@ def get_type_qualifiers() -> tuple[str, TypeEngine[Any], str]:
             return "".join(pre), column_type, "]" * post_size
 
         def render_python_type(column_type: TypeEngine[Any]) -> str:
+            # Check if this is an enum column with a Python enum class
+            if isinstance(column_type, Enum):
+                table_name = column.table.name
+                column_name = column.name
+                if (table_name, column_name) in self.enum_classes:
+                    enum_class_name = self.enum_classes[(table_name, column_name)]
+                    return enum_class_name
+
             if isinstance(column_type, DOMAIN):
                 column_type = column_type.data_type
 

src/sqlacodegen/utils.py

@@ -210,10 +210,4 @@ def decode_postgresql_sequence(clause: TextClause) -> tuple[str | None, str | No
 
 
 def get_stdlib_module_names() -> set[str]:
-    major, minor = sys.version_info.major, sys.version_info.minor
-    if (major, minor) > (3, 9):
-        return set(sys.builtin_module_names) | set(sys.stdlib_module_names)
-    else:
-        from stdlib_list import stdlib_list
-
-        return set(sys.builtin_module_names) | set(stdlib_list(f"{major}.{minor}"))
+    return set(sys.builtin_module_names) | set(sys.stdlib_module_names)