Add support to postgres range fields

CHANGELOG.md

@@ -8,10 +8,12 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 ## [Unreleased](https://github.com/model-bakers/model_bakery/tree/master)
 
 ### Added
-- Add isort and fix imports [PR #77](https://github.com/model-bakers/model_bakery/pull/77)
+- Support to Postgres fields: `DecimalRangeField`, `FloatRangeField`, `IntegerRangeField`, `BigIntegerRangeField`, `DateRangeField`, `DateTimeRangeField` [PR #80](https://github.com/model-bakers/model_bakery/pull/80)
 
 ### Changed
+- Add isort and fix imports [PR #77](https://github.com/model-bakers/model_bakery/pull/77)
 - Enable `seq` to be imported from `baker` [PR #76](https://github.com/model-bakers/model_bakery/pull/76)
+- Fix PostGIS model registration [PR #67](https://github.com/model-bakers/model_bakery/pull/67)
 
 ### Removed
 

dev_requirements.txt

@@ -8,3 +8,4 @@ sphinx
 tox
 twine==2.0.0
 wheel
+psycopg2==2.8.5

docs/source/how_bakery_behaves.rst

@@ -49,6 +49,7 @@ Currently supported fields
 * ``FileField``, ``ImageField``
 * ``JSONField``, ``ArrayField``, ``HStoreField``
 * ``CICharField``, ``CIEmailField``, ``CITextField``
+* ``DecimalRangeField``, ``IntegerRangeField``, ``BigIntegerRangeField``, ``DateRangeField``, ``DateTimeRangeField``
 
 Require ``django.contrib.gis`` in ``INSTALLED_APPS``:
 

model_bakery/baker.py

@@ -27,10 +27,10 @@
     ModelNotFound,
     RecipeIteratorEmpty,
 )
-from .utils import import_from_str
 
 # Enable seq to be imported from baker
 from .utils import seq  # NoQA
+from .utils import import_from_str
 
 recipes = None
 

model_bakery/generators.py

@@ -1,3 +1,5 @@
+from decimal import Decimal
+
 from django.db.backends.base.operations import BaseDatabaseOperations
 from django.db.models import (
     BigIntegerField,
@@ -70,6 +72,25 @@
     CIEmailField = None
     CITextField = None
 
+try:
+    from django.contrib.postgres.fields.ranges import DecimalRangeField
+except ImportError:
+    DecimalRangeField = None
+try:
+    from django.contrib.postgres.fields.ranges import (
+        BigIntegerRangeField,
+        DateRangeField,
+        DateTimeRangeField,
+        FloatRangeField,
+        IntegerRangeField,
+    )
+except ImportError:
+    IntegerRangeField = None
+    BigIntegerRangeField = None
+    FloatRangeField = None
+    DateRangeField = None
+    DateTimeRangeField = None
+
 
 def _make_integer_gen_by_range(field_type):
     min_int, max_int = BaseDatabaseOperations.integer_field_ranges[field_type.__name__]
@@ -132,13 +153,26 @@ def gen_integer():
     default_mapping[PositiveBigIntegerField] = _make_integer_gen_by_range(
         PositiveBigIntegerField
     )
+if DecimalRangeField:
+    default_mapping[DecimalRangeField] = random_gen.gen_pg_numbers_range(Decimal)
+if IntegerRangeField:
+    default_mapping[IntegerRangeField] = random_gen.gen_pg_numbers_range(int)
+if BigIntegerRangeField:
+    default_mapping[BigIntegerRangeField] = random_gen.gen_pg_numbers_range(int)
+if FloatRangeField:
+    default_mapping[FloatRangeField] = random_gen.gen_pg_numbers_range(float)
+if DateRangeField:
+    default_mapping[DateRangeField] = random_gen.gen_date_range
+if DateTimeRangeField:
+    default_mapping[DateTimeRangeField] = random_gen.gen_datetime_range
 
 
 # Add GIS fields
 
 
 def get_type_mapping():
     from django.contrib.contenttypes.models import ContentType
+
     from .gis import default_gis_mapping
 
     mapping = default_mapping.copy()

model_bakery/gis.py

@@ -9,18 +9,19 @@
 __all__ = ["BAKER_GIS", "default_gis_mapping"]
 
 if BAKER_GIS:
-    from . import random_gen
     from django.contrib.gis.db.models import (
+        GeometryCollectionField,
         GeometryField,
-        PointField,
         LineStringField,
-        PolygonField,
-        MultiPointField,
         MultiLineStringField,
+        MultiPointField,
         MultiPolygonField,
-        GeometryCollectionField,
+        PointField,
+        PolygonField,
     )
 
+    from . import random_gen
+
     default_gis_mapping[GeometryField] = random_gen.gen_geometry
     default_gis_mapping[PointField] = random_gen.gen_point
     default_gis_mapping[LineStringField] = random_gen.gen_line_string

model_bakery/random_gen.py

@@ -200,8 +200,8 @@ def gen_interval(interval_key="milliseconds"):
 
 
 def gen_content_type():
-    from django.contrib.contenttypes.models import ContentType
     from django.apps import apps
+    from django.contrib.contenttypes.models import ContentType
 
     try:
         return ContentType.objects.get_for_model(choice(apps.get_models()))
@@ -252,7 +252,7 @@ def gen_related(model, **attrs):
 
 
 def gen_m2m(model, **attrs):
-    from .baker import make, MAX_MANY_QUANTITY
+    from .baker import MAX_MANY_QUANTITY, make
 
     return make(model, _quantity=MAX_MANY_QUANTITY, **attrs)
 
@@ -310,3 +310,31 @@ def gen_geometry():
 
 def gen_geometry_collection():
     return "GEOMETRYCOLLECTION ({})".format(gen_point(),)
+
+
+def gen_pg_numbers_range(number_cast=int):
+    def gen_range():
+        from psycopg2._range import NumericRange
+
+        base_num = gen_integer(1, 100000)
+        return NumericRange(number_cast(-1 * base_num), number_cast(base_num))
+
+    return gen_range
+
+
+def gen_date_range():
+    from psycopg2.extras import DateRange
+
+    base_date = gen_date()
+    interval = gen_interval()
+    args = sorted([base_date - interval, base_date + interval])
+    return DateRange(*args)
+
+
+def gen_datetime_range():
+    from psycopg2.extras import DateTimeTZRange
+
+    base_datetime = gen_datetime()
+    interval = gen_interval()
+    args = sorted([base_datetime - interval, base_datetime + interval])
+    return DateTimeTZRange(*args)

tests/generic/models.py

@@ -94,17 +94,37 @@ class Person(models.Model):
             CIEmailField,
             CITextField,
         )
+        from django.contrib.postgres.fields.ranges import (
+            IntegerRangeField,
+            BigIntegerRangeField,
+            FloatRangeField,
+            DateRangeField,
+            DateTimeRangeField,
+        )
 
         acquaintances = ArrayField(models.IntegerField())
         data = JSONField()
         hstore_data = HStoreField()
         ci_char = CICharField(max_length=30)
         ci_email = CIEmailField()
         ci_text = CITextField()
+        int_range = IntegerRangeField()
+        bigint_range = BigIntegerRangeField()
+        float_range = FloatRangeField()
+        date_range = DateRangeField()
+        datetime_range = DateTimeRangeField()
     except ImportError:
         # Skip PostgreSQL-related fields
         pass
 
+    try:
+        from django.contrib.postgres.fields.ranges import DecimalRangeField
+
+        decimal_range = DecimalRangeField()
+    except ImportError:
+        # Django version lower than 2.2
+        pass
+
     if BAKER_GIS:
         geom = models.GeometryField()
         point = models.PointField()

tests/test_baker.py

@@ -148,7 +148,7 @@ def test_accepts_generators_with_quantity(self):
             _quantity=5,
         )
         assert models.Person.objects.count() == 5
-        p1, p2, p3, p4, p5 = models.Person.objects.all()
+        p1, p2, p3, p4, p5 = models.Person.objects.all().order_by("pk")
         assert "a" == p1.name
         assert "d1" == p1.id_document
         assert "b" == p2.name

tests/test_filling_fields.py

@@ -28,13 +28,30 @@
         HStoreField,
         JSONField,
     )
+    from django.contrib.postgres.fields.ranges import (
+        IntegerRangeField,
+        BigIntegerRangeField,
+        FloatRangeField,
+        DateRangeField,
+        DateTimeRangeField,
+    )
 except ImportError:
     ArrayField = None
     JSONField = None
     HStoreField = None
     CICharField = None
     CIEmailField = None
     CITextField = None
+    IntegerRangeField = None
+    BigIntegerRangeField = None
+    FloatRangeField = None
+    DateRangeField = None
+    DateTimeRangeField = None
+
+try:
+    from django.contrib.postgres.fields.ranges import DecimalRangeField
+except ImportError:
+    DecimalRangeField = None
 
 
 @pytest.fixture
@@ -389,6 +406,70 @@ def test_filling_citextfield(self, person):
         assert isinstance(ci_text_field, CITextField)
         assert isinstance(person.ci_text, str)
 
+    @pytest.mark.skipif(
+        not DecimalRangeField,
+        reason="Django version does not support DecimalRangeField",
+    )
+    def test_filling_decimal_range_field(self, person):
+        from psycopg2._range import NumericRange
+
+        decimal_range_field = models.Person._meta.get_field("decimal_range")
+        assert isinstance(decimal_range_field, DecimalRangeField)
+        assert isinstance(person.decimal_range, NumericRange)
+        assert isinstance(person.decimal_range.lower, Decimal)
+        assert isinstance(person.decimal_range.upper, Decimal)
+        assert person.decimal_range.lower < person.decimal_range.upper
+
+    def test_filling_integer_range_field(self, person):
+        from psycopg2._range import NumericRange
+
+        int_range_field = models.Person._meta.get_field("int_range")
+        assert isinstance(int_range_field, IntegerRangeField)
+        assert isinstance(person.int_range, NumericRange)
+        assert isinstance(person.int_range.lower, int)
+        assert isinstance(person.int_range.upper, int)
+        assert person.int_range.lower < person.int_range.upper
+
+    def test_filling_integer_range_field_for_big_int(self, person):
+        from psycopg2._range import NumericRange
+
+        bigint_range_field = models.Person._meta.get_field("bigint_range")
+        assert isinstance(bigint_range_field, BigIntegerRangeField)
+        assert isinstance(person.bigint_range, NumericRange)
+        assert isinstance(person.bigint_range.lower, int)
+        assert isinstance(person.bigint_range.upper, int)
+        assert person.bigint_range.lower < person.bigint_range.upper
+
+    def test_filling_float_range_field(self, person):
+        from psycopg2._range import NumericRange
+
+        float_range_field = models.Person._meta.get_field("float_range")
+        assert isinstance(float_range_field, FloatRangeField)
+        assert isinstance(person.float_range, NumericRange)
+        assert isinstance(person.float_range.lower, float)
+        assert isinstance(person.float_range.upper, float)
+        assert person.float_range.lower < person.float_range.upper
+
+    def test_filling_date_range_field(self, person):
+        from psycopg2.extras import DateRange
+
+        date_range_field = models.Person._meta.get_field("date_range")
+        assert isinstance(date_range_field, DateRangeField)
+        assert isinstance(person.date_range, DateRange)
+        assert isinstance(person.date_range.lower, date)
+        assert isinstance(person.date_range.upper, date)
+        assert person.date_range.lower < person.date_range.upper
+
+    def test_filling_datetime_range_field(self, person):
+        from psycopg2.extras import DateTimeTZRange
+
+        datetime_range_field = models.Person._meta.get_field("datetime_range")
+        assert isinstance(datetime_range_field, DateTimeRangeField)
+        assert isinstance(person.datetime_range, DateTimeTZRange)
+        assert isinstance(person.datetime_range.lower, datetime)
+        assert isinstance(person.datetime_range.upper, datetime)
+        assert person.datetime_range.lower < person.datetime_range.upper
+
 
 @pytest.mark.skipif(
     connection.vendor != "postgresql", reason="PostgreSQL specific tests"

tox.ini

@@ -9,7 +9,7 @@ envlist =
 [travis]
 python =
     3.5: py35
-    3.6: py36,flake8,isort,black
+    3.6: py36,flake8,black
     3.7: py37
     3.8: py38
 
@@ -37,11 +37,6 @@ deps=flake8
 basepython=python3
 commands=flake8 model_bakery {posargs}
 
-[testenv:isort]
-deps=isort
-basepython=python3
-commands=isort model_bakery {posargs}
-
 [testenv:black]
 deps=black
 basepython=python3