#3287: settings.py, ensure python 2+3 string compatibility

jeff1evesque · Oct 9, 2018 · 64d5c9e · 64d5c9e
1 parent 5c914f3
commit 64d5c9e
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/brain/validator/settings.py b/brain/validator/settings.py
@@ -9,6 +9,7 @@
 from flask import current_app
 from voluptuous import Schema, Required, Optional, All, Any, Coerce, In, Length
 from voluptuous.humanize import validate_with_humanized_errors
+from six import string_types
 
 
 class Validator(object):
@@ -50,17 +51,17 @@ def validate_settings(self, premodel_settings, session_type):
         if session_type in ['data_new', 'data_append']:
             if premodel_settings['stream'] == 'True' or session_type == 'data_new':
                 schema = Schema({
-                    Required('collection'): All(unicode, Length(min=1)),
+                    Required('collection'): All(string_types, Length(min=1)),
                     Required('dataset_type'): In(dataset_type),
                     Required('model_type'): In(model_type),
                     Required('session_type'): Any('data_new', 'data_append'),
-                    Required('session_name'): All(unicode, Length(min=1)),
+                    Required('session_name'): All(string_types, Length(min=1)),
                     Optional('stream'): Any('True', 'False'),
                 })
 
             else:
                 schema = Schema({
-                    Required('collection'): All(unicode, Length(min=1)),
+                    Required('collection'): All(string_types, Length(min=1)),
                     Required('dataset_type'): In(dataset_type),
                     Required('model_type'): In(model_type),
                     Required('session_type'): Any('data_new', 'data_append'),
@@ -70,7 +71,7 @@ def validate_settings(self, premodel_settings, session_type):
         # validation on 'model_generate' session
         if session_type == 'model_generate':
             schema = Schema({
-                Required('collection'): All(unicode, Length(min=1)),
+                Required('collection'): All(string_types, Length(min=1)),
                 Required('model_type'): In(model_type),
                 Required('session_type'): 'model_generate',
                 Optional('stream'): Any('True', 'False'),
@@ -82,7 +83,7 @@ def validate_settings(self, premodel_settings, session_type):
         # validation on 'model_predict' session
         elif session_type == 'model_predict':
             schema = Schema({
-                Required('collection'): All(unicode, Length(min=1)),
+                Required('collection'): All(string_types, Length(min=1)),
                 Optional('stream'): Any('True', 'False'),
                 Required('prediction_input[]'): [
                     Any(Coerce(int), Coerce(float)),