Refactor function which chooses column providers

Author: Iain-S

sqlsynthgen/make.py

@@ -185,7 +185,7 @@ def _get_default_generator(tables_module: ModuleType, column: Any) -> RowGenerat
             variable_names,
             generator_function,
             generator_arguments,
-        ) = _get_mimesis_function_for_colum(column)
+        ) = _get_provider_for_column(column)
 
     return RowGenerator(
         primary_key=column.primary_key,
@@ -196,7 +196,7 @@ def _get_default_generator(tables_module: ModuleType, column: Any) -> RowGenerat
     )
 
 
-def _get_mimesis_function_for_colum(column: Any) -> Tuple[List[str], str, List[str]]:
+def _get_provider_for_column(column: Any) -> Tuple[List[str], str, List[str]]:
     """
     Get a default Mimesis provider and its arguments for a SQL column type.
 
@@ -209,37 +209,35 @@ def _get_mimesis_function_for_colum(column: Any) -> Tuple[List[str], str, List[s
     """
     variable_names: List[str] = [column.name]
     generator_arguments: List[str] = []
-    generator_function: str = ""
 
     column_type = type(column.type)
     column_size: Optional[int] = getattr(column.type, "length", None)
 
-    # ToDo Add tests and then add issubclass for all of these
-    # sqlalchemy.dialects.mysql.types.INTEGER
-    if column_type == sqltypes.BigInteger:
-        generator_function = "generic.numeric.integer_number"
-    if column_type == sqltypes.Integer or issubclass(column_type, sqltypes.Integer):
-        generator_function = "generic.numeric.integer_number"
-    elif column_type == sqltypes.Boolean:
-        generator_function = "generic.development.boolean"
-    elif column_type == sqltypes.Date:
-        generator_function = "generic.datetime.date"
-    elif column_type == sqltypes.DateTime:
-        generator_function = "generic.datetime.datetime"
-    elif column_type in {sqltypes.Float, sqltypes.Numeric}:
-        generator_function = "generic.numeric.float_number"
-    elif column_type == sqltypes.Integer:
-        generator_function = "generic.numeric.integer_number"
-    elif column_type == sqltypes.LargeBinary:
-        generator_function = "generic.bytes_provider.bytes"
-    elif column_type in {sqltypes.String, sqltypes.Text} and column_size is None:
-        generator_function = "generic.text.color"
-    elif column_type in {sqltypes.String, sqltypes.Text} and column_size is not None:
-        generator_function = "generic.person.password"
-        generator_arguments.append(str(column_size))
-    else:
+    mapping = {
+        (sqltypes.Integer, False): "generic.numeric.integer_number",
+        (sqltypes.Boolean, False): "generic.development.boolean",
+        (sqltypes.Date, False): "generic.datetime.date",
+        (sqltypes.DateTime, False): "generic.datetime.datetime",
+        (sqltypes.Numeric, False): "generic.numeric.float_number",
+        (sqltypes.LargeBinary, False): "generic.bytes_provider.bytes",
+        (sqltypes.String, False): "generic.text.color",
+        (sqltypes.String, True): "generic.person.password",
+    }
+
+    generator_function = mapping.get((column_type, column_size is not None), None)
+
+    if not generator_function:
+        for key, value in mapping.items():
+            if issubclass(column_type, key[0]) and key[1] == (column_size is not None):
+                generator_function = value
+                break
+
+    if not generator_function:
         raise ValueError(f"Unsupported SQLAlchemy type: {column_type}")
 
+    if column_size:
+        generator_arguments.append(str(column_size))
+
     return variable_names, generator_function, generator_arguments
 
 

tests/test_make.py

@@ -9,8 +9,15 @@
 import yaml
 from pydantic import PostgresDsn
 from pydantic.tools import parse_obj_as
-
-from sqlsynthgen.make import make_src_stats, make_table_generators, make_tables_file
+from sqlalchemy import BigInteger, Column, String
+from sqlalchemy.dialects.mysql.types import INTEGER
+
+from sqlsynthgen.make import (
+    _get_provider_for_column,
+    make_src_stats,
+    make_table_generators,
+    make_tables_file,
+)
 from tests.examples import example_orm
 from tests.utils import RequiresDBTestCase, SSGTestCase, get_test_settings
 
@@ -119,6 +126,54 @@ def test_make_generators_force_overwrite(
 
         self.assertEqual(expected, actual)
 
+    def test__get_provider_for_column(self) -> None:
+        """Test the _get_provider_for_column function."""
+
+        # Simple case
+        (
+            variable_name,
+            generator_function,
+            generator_arguments,
+        ) = _get_provider_for_column(Column("myint", BigInteger))
+        self.assertListEqual(
+            variable_name,
+            ["myint"],
+        )
+        self.assertEqual(
+            generator_function,
+            "generic.numeric.integer_number",
+        )
+        self.assertEqual(
+            generator_arguments,
+            [],
+        )
+
+        # Column type from another dialect
+        _, generator_function, __ = _get_provider_for_column(Column("myint", INTEGER))
+        self.assertEqual(
+            generator_function,
+            "generic.numeric.integer_number",
+        )
+
+        # Text value with length
+        (
+            variable_name,
+            generator_function,
+            generator_arguments,
+        ) = _get_provider_for_column(Column("mystring", String(100)))
+        self.assertEqual(
+            variable_name,
+            ["mystring"],
+        )
+        self.assertEqual(
+            generator_function,
+            "generic.person.password",
+        )
+        self.assertEqual(
+            generator_arguments,
+            ["100"],
+        )
+
 
 class TestMakeTables(SSGTestCase):
     """Test the make_tables function."""