import Union

src/sqlacodegen/generators.py

@@ -1215,11 +1215,15 @@ def render_column_attribute(self, column_attr: ColumnAttribute) -> str:
             return f"{column_attr.name} = {rendered_column}"
         else:
             try:
-                # FIXME: this is due to sqlalchemy associate JSON to dict, but actually there is a case that JSON is list
-                # https://github.com/sqlalchemy/sqlalchemy/blob/main/lib/sqlalchemy/sql/sqltypes.py#L2665
-                python_type = column.type.python_type if not isinstance(column.type, JSONB) else Union[dict, list]
+                python_type = column.type.python_type
                 python_type_name = python_type.__name__
-                if python_type.__module__ == "builtins":
+                if isinstance(column.type, JSONB):
+                    # FIXME: this is due to sqlalchemy associate JSON to dict, but actually there is a case that JSON is list
+                    # https://github.com/sqlalchemy/sqlalchemy/blob/main/lib/sqlalchemy/sql/sqltypes.py#L2665
+                    self.add_literal_import("typing", "Union")
+                    python_type = Union[dict, list]
+                    column_python_type = "Union[dict, list]"
+                elif python_type.__module__ == "builtins":
                     column_python_type = python_type_name
                 else:
                     python_type_module = python_type.__module__