Refactor ModelBuilder and RandomBuilder

piccolo/testing/model_builder.py

@@ -1,31 +1,15 @@
 from __future__ import annotations
 
-import datetime
 import json
 import typing as t
-from decimal import Decimal
-from uuid import UUID
 
-from piccolo.columns import JSON, JSONB, Array, Column, ForeignKey
+from piccolo.columns import JSON, JSONB, Column, ForeignKey
 from piccolo.custom_types import TableInstance
 from piccolo.testing.random_builder import RandomBuilder
 from piccolo.utils.sync import run_sync
 
 
 class ModelBuilder:
-    __DEFAULT_MAPPER: t.Dict[t.Type, t.Callable] = {
-        bool: RandomBuilder.next_bool,
-        bytes: RandomBuilder.next_bytes,
-        datetime.date: RandomBuilder.next_date,
-        datetime.datetime: RandomBuilder.next_datetime,
-        float: RandomBuilder.next_float,
-        int: RandomBuilder.next_int,
-        str: RandomBuilder.next_str,
-        datetime.time: RandomBuilder.next_time,
-        datetime.timedelta: RandomBuilder.next_timedelta,
-        UUID: RandomBuilder.next_uuid,
-    }
-
     @classmethod
     async def build(
         cls,
@@ -159,29 +143,7 @@ def _randomize_attribute(cls, column: Column) -> t.Any:
             Column class to randomize.
 
         """
-        random_value: t.Any
-        if column.value_type == Decimal:
-            precision, scale = column._meta.params["digits"] or (4, 2)
-            random_value = RandomBuilder.next_float(
-                maximum=10 ** (precision - scale), scale=scale
-            )
-        elif column.value_type == datetime.datetime:
-            tz_aware = getattr(column, "tz_aware", False)
-            random_value = RandomBuilder.next_datetime(tz_aware=tz_aware)
-        elif column.value_type == list:
-            length = RandomBuilder.next_int(maximum=10)
-            base_type = t.cast(Array, column).base_column.value_type
-            random_value = [
-                cls.__DEFAULT_MAPPER[base_type]() for _ in range(length)
-            ]
-        elif column._meta.choices:
-            random_value = RandomBuilder.next_enum(column._meta.choices)
-        else:
-            random_value = cls.__DEFAULT_MAPPER[column.value_type]()
-
-        if "length" in column._meta.params and isinstance(random_value, str):
-            return random_value[: column._meta.params["length"]]
-        elif isinstance(column, (JSON, JSONB)):
+        random_value: t.Any = RandomBuilder._build(column)
+        if isinstance(column, (JSON, JSONB)):
             return json.dumps({"value": random_value})
-
         return random_value

piccolo/testing/random_builder.py

@@ -1,12 +1,50 @@
+from __future__ import annotations
+
 import datetime
 import enum
 import random
 import string
 import typing as t
 import uuid
+from decimal import Decimal
+from functools import partial
+from uuid import UUID
+
+from piccolo.columns import Array, Column
 
 
 class RandomBuilder:
+    @classmethod
+    def _build(cls, column: Column) -> t.Any:
+        if e := column._meta.choices:
+            return cls.next_enum(e)
+
+        mapper: t.Dict[t.Type, t.Callable] = {
+            bool: cls.next_bool,
+            bytes: cls.next_bytes,
+            datetime.date: cls.next_date,
+            datetime.datetime: partial(
+                cls.next_datetime, getattr(column, "tz_aware", False)
+            ),
+            float: cls.next_float,
+            Decimal: partial(
+                cls.next_decimal, column._meta.params.get("digits")
+            ),
+            int: cls.next_int,
+            str: partial(cls.next_str, column._meta.params.get("length")),
+            datetime.time: cls.next_time,
+            datetime.timedelta: cls.next_timedelta,
+            UUID: cls.next_uuid,
+        }
+
+        random_value_callable = mapper.get(column.value_type)
+        if random_value_callable is None:
+            random_value_callable = partial(
+                cls.next_list,
+                mapper[t.cast(Array, column).base_column.value_type],
+            )
+        return random_value_callable()
+
     @classmethod
     def next_bool(cls) -> bool:
         return random.choice([True, False])
@@ -43,12 +81,18 @@ def next_enum(cls, e: t.Type[enum.Enum]) -> t.Any:
     def next_float(cls, minimum=0, maximum=2147483647, scale=5) -> float:
         return round(random.uniform(minimum, maximum), scale)
 
+    @classmethod
+    def next_decimal(cls, digits: t.Tuple[int, int] | None = (4, 2)) -> float:
+        precision, scale = digits or (4, 2)
+        return cls.next_float(maximum=10 ** (precision - scale), scale=scale)
+
     @classmethod
     def next_int(cls, minimum=0, maximum=2147483647) -> int:
         return random.randint(minimum, maximum)
 
     @classmethod
-    def next_str(cls, length=16) -> str:
+    def next_str(cls, length: int | None = 16) -> str:
+        length = length or 16
         return "".join(
             random.choice(string.ascii_letters) for _ in range(length)
         )
@@ -72,3 +116,8 @@ def next_timedelta(cls) -> datetime.timedelta:
     @classmethod
     def next_uuid(cls) -> uuid.UUID:
         return uuid.uuid4()
+
+    @classmethod
+    def next_list(cls, callable_: t.Callable) -> t.List[t.Any]:
+        length = cls.next_int(maximum=10)
+        return [callable_() for _ in range(length)]

tests/testing/test_random_builder.py

@@ -1,5 +1,7 @@
+import datetime
 import unittest
 from enum import Enum
+from uuid import UUID
 
 from piccolo.testing.random_builder import RandomBuilder
 
@@ -31,6 +33,10 @@ class Color(Enum):
         random_enum = RandomBuilder.next_enum(Color)
         self.assertIsInstance(random_enum, int)
 
+    def test_next_decimal(self):
+        random_decimal = RandomBuilder.next_decimal((5, 2))
+        self.assertLessEqual(random_decimal, 1000)
+
     def test_next_float(self):
         random_float = RandomBuilder.next_float(maximum=1000)
         self.assertLessEqual(random_float, 1000)
@@ -52,3 +58,24 @@ def test_next_timedelta(self):
 
     def test_next_uuid(self):
         RandomBuilder.next_uuid()
+
+    def test_next_list(self):
+        # `RandomBuilder.next_decimal` will return `float`
+        reversed_mapper = {
+            RandomBuilder.next_bool: bool,
+            RandomBuilder.next_bytes: bytes,
+            RandomBuilder.next_date: datetime.date,
+            RandomBuilder.next_datetime: datetime.datetime,
+            RandomBuilder.next_float: float,
+            RandomBuilder.next_decimal: float,
+            RandomBuilder.next_int: int,
+            RandomBuilder.next_str: str,
+            RandomBuilder.next_time: datetime.time,
+            RandomBuilder.next_timedelta: datetime.timedelta,
+            RandomBuilder.next_uuid: UUID,
+        }
+
+        for callable_, typ in reversed_mapper.items():
+            random_list = RandomBuilder.next_list(callable_)
+            self.assertIsInstance(random_list, list)
+            self.assertTrue(all(isinstance(elem, typ) for elem in random_list))