From 603d42792d460a6a2cae4f7e0b8d32b8b21e3cdb Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Fri, 15 Nov 2019 22:00:21 +0200
Subject: [PATCH] Fix jsonschema deprecation warning

---
 trains/backend_api/session/datamodel.py | 33 +++++++++++++++++++++----
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/trains/backend_api/session/datamodel.py b/trains/backend_api/session/datamodel.py
index c3a0f896..2294e774 100644
--- a/trains/backend_api/session/datamodel.py
+++ b/trains/backend_api/session/datamodel.py
@@ -31,6 +31,22 @@ def schema_property(name):
     return init
 
 
+# Support both jsonschema >= 3.0.0 and <= 2.6.0
+_CustomValidator = None
+try:
+    from jsonschema import TypeChecker, Draft7Validator
+
+    def _is_array(checker, instance):
+        return isinstance(instance, (list, tuple))
+
+    _CustomValidator = jsonschema.validators.extend(
+        Draft7Validator,
+        type_checker=Draft7Validator.TYPE_CHECKER.redefine("array", _is_array)
+    )
+except ImportError:
+    pass
+
+
 class DataModel(object):
     """ Data Model"""
     _schema = None
@@ -66,11 +82,18 @@ class DataModel(object):
         }
 
     def validate(self, schema=None):
-        jsonschema.validate(
-            self.to_dict(),
-            schema or self._schema,
-            types=dict(array=(list, tuple), integer=six.integer_types),
-        )
+        if _CustomValidator is None:
+            jsonschema.validate(
+                self.to_dict(),
+                schema or self._schema,
+                types=dict(array=(list, tuple), integer=six.integer_types),
+            )
+        else:
+            jsonschema.validate(
+                self.to_dict(),
+                schema or self._schema,
+                cls=_CustomValidator,
+            )
 
     def __repr__(self):
         return '<{}.{}: {}>'.format(