Better type checks

Fixes #170 #171

4 files changed, 28 insertions, 5 deletions

diff --git a/docs/CHANGES.rst b/docs/CHANGES.rst
index 7373308..ca1fdd6 100644
--- a/docs/CHANGES.rst
+++ b/docs/CHANGES.rst
@@ -4,7 +4,7 @@ Changelog
 1.0.14 (unreleased)
 -------------------
 
-- Nothing changed yet.
+- better type checks [kiorky]
 
 
 1.0.13 (2021-05-06)
diff --git a/src/croniter/__init__.py b/src/croniter/__init__.py
index f9f62f8..406263b 100644
--- a/src/croniter/__init__.py
+++ b/src/croniter/__init__.py
@@ -3,6 +3,7 @@ from __future__ import absolute_import
 from .croniter import (
     croniter,
     croniter_range,
+    CroniterBadTypeRangeError,  # noqa
     CroniterBadDateError,  # noqa
     CroniterBadCronError,  # noqa
     CroniterNotAlphaError, # noqa
diff --git a/src/croniter/croniter.py b/src/croniter/croniter.py
index 52617ed..e96ea52 100644
--- a/src/croniter/croniter.py
+++ b/src/croniter/croniter.py
@@ -48,6 +48,10 @@ class CroniterError(ValueError):
     pass
 
 
+class CroniterBadTypeRangeError(TypeError):
+    """."""
+
+
 class CroniterBadCronError(CroniterError):
     """ Syntax, unknown value, or range error within a cron expression """
     pass
@@ -795,9 +799,15 @@ def croniter_range(start, stop, expr_format, ret_type=None, day_or=True, exclude
     """
     _croniter = _croniter or croniter
     auto_rt = datetime.datetime
-    if type(start) != type(stop):
-        raise TypeError("The start and stop must be same type.  {0} != {1}".
-                        format(type(start), type(stop)))
+    # type is used in first if branch for perfs reasons
+    if (
+        type(start) != type(stop) and not (
+            isinstance(start, type(stop)) or
+            isinstance(stop, type(start)))
+    ):
+        raise CroniterBadTypeRangeError(
+            "The start and stop must be same type.  {0} != {1}".
+            format(type(start), type(stop)))
     if isinstance(start, (float, int)):
         start, stop = (datetime.datetime.utcfromtimestamp(t) for t in (start, stop))
         auto_rt = float
diff --git a/src/croniter/tests/test_croniter_range.py b/src/croniter/tests/test_croniter_range.py
index b9cace0..29fc0d0 100755
--- a/src/croniter/tests/test_croniter_range.py
+++ b/src/croniter/tests/test_croniter_range.py
@@ -4,10 +4,14 @@
 import unittest
 from datetime import datetime, timedelta
 import pytz
-from croniter import croniter, croniter_range, CroniterBadDateError, CroniterBadCronError
+from croniter import croniter, croniter_range, CroniterBadDateError, CroniterBadCronError, CroniterBadTypeRangeError
 from croniter.tests import base
 
 
+class mydatetime(datetime):
+    """."""
+
+
 class CroniterRangeTest(base.TestCase):
 
     def test_1day_step(self):
@@ -148,6 +152,14 @@ class CroniterRangeTest(base.TestCase):
             i = croniter_range(datetime(2020, 1, 1), datetime(2020, 12, 31), cron, _croniter=croniter_nosec)
             next(i)
 
+    def test_dt_types(self):
+        start = mydatetime(2020, 9, 24)
+        stop = datetime(2020, 9, 28)
+        try:
+            list(croniter_range(start, stop, '0 0 * * *'))
+        except CroniterBadTypeRangeError:
+            self.fail('should not be triggered')
+
 
 if __name__ == '__main__':
     unittest.main()