Allow Mockgun to apply schema modifications

shotgun_api3/lib/mockgun/mockgun.py

@@ -114,6 +114,7 @@
 
 """
 
+import copy
 import datetime
 
 from ... import ShotgunError
@@ -227,13 +228,48 @@ def schema_read(self):
         return self._schema
 
     def schema_field_create(self, entity_type, data_type, display_name, properties=None):
-        raise NotImplementedError
+        _properties = {
+            "data_type": {"editable": True, "value": None},
+            "description": {"editable": True, "value": ""},
+            "editable": {"editable": True, "value": True},
+            "entity_type": {"editable": True, "value": None},
+            "mandatory": {"editable": True, "value": False},
+            "name": {"editable": True, "value": None},
+            "properties": {
+                "default_value": {"editable": True, "value": None},
+                "regex_validation": {"editable": True, "value": ""},
+                "regex_validation_enabled": {"editable": True, "value": None},
+                "summary_default": {"editable": True, "value": "none"}
+            },
+            "ui_value_displayable": {"editable": True, "value": True},
+            "unique": {"editable": True, "value": True},
+            "visible": {"editable": True, "value": True}
+        }
+        _properties["entity_type"]["value"] = entity_type
+        _properties["entity_type"]["editable"] = False
+
+        _properties["data_type"]["value"] = data_type
+        _properties["data_type"]["editable"] = False
+
+        _properties["name"]["value"] = display_name
+        _properties["name"]["editable"] = False
+
+        if isinstance(properties, dict) and properties:
+            for prop in properties:
+                self._set_property(_properties, prop, properties[prop])
+
+        self._schema[entity_type][display_name] = _properties
 
     def schema_field_update(self, entity_type, field_name, properties):
-        raise NotImplementedError
+        _properties = copy.deepcopy(self._schema[entity_type][field_name])
+        for key in properties:
+            self._set_property(_properties, key, properties[key])
+        self._schema[entity_type][field_name] = _properties
 
     def schema_field_delete(self, entity_type, field_name):
-        raise NotImplementedError
+        if field_name not in self._schema[entity_type]:
+            return
+        del self._schema[entity_type][field_name]
 
     def schema_entity_read(self):
         return self._schema_entity
@@ -422,8 +458,29 @@ def upload(self, entity_type, entity_id, path, field_name=None, display_name=Non
     def upload_thumbnail(self, entity_type, entity_id, path, **kwargs):
         pass
 
+    def dump_schemas(self):
+        schema_path, schema_entity_path = self.get_schema_paths()
+        SchemaFactory.set_schemas(self.schema_read(), schema_path, self.schema_entity_read(), schema_entity_path)
+
     ###################################################################################################
     # internal methods and members
+    @classmethod
+    def _set_property(cls, property_data, property_key, property_value):
+        result = False
+        for key in property_data:
+            if key == property_key:
+                property_data[key]["value"] = property_value
+                result = True
+                break
+            elif isinstance(property_data[key], dict):
+                _result = cls._set_property(property_data[key], property_key, property_value)
+                if _result is True:
+                    result = _result
+                    break
+        if result is False and "properties" in property_data:
+            property_data["properties"][property_key] = {"editable": True, "value": property_value}
+            result = True
+        return result
 
     def _validate_entity_type(self, entity_type):
         if entity_type not in self._schema:

shotgun_api3/lib/mockgun/schema.py

@@ -79,6 +79,27 @@ def get_schemas(cls, schema_path, schema_entity_path):
 
         return cls._schema_cache, cls._schema_entity_cache
 
+    @classmethod
+    def set_schemas(cls, schema, schema_path, schema_entity, schema_entity_path):
+        """
+        Retrieves the schemas from disk.
+
+        :param str schema: Schema data.
+        :param str schema_path: Path to the schema on disk.
+        :param str schema_entity: Entities schema data.
+        :param str schema_entity_path: Path to the entities schema on disk.
+        """
+
+        if schema_path != cls._schema_cache_path:
+            cls._schema_cache_path = schema_path
+        cls._schema_cache = schema
+        cls._write_file(schema, schema_path)
+
+        if schema_entity_path != cls._schema_entity_cache_path:
+            cls._schema_entity_cache_path = schema_entity_path
+        cls._schema_entity_cache = schema_entity
+        cls._write_file(schema_entity, schema_entity_path)
+
     @classmethod
     def _read_file(cls, path):
         fh = open(path, "rb")
@@ -87,6 +108,14 @@ def _read_file(cls, path):
         finally:
             fh.close()
 
+    @classmethod
+    def _write_file(cls, data, path):
+        fh = open(path, "wb")
+        try:
+            return pickle.dump(data, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL)
+        finally:
+            fh.close()
+
 
 # Highest protocol that Python 2.4 supports, which is the earliest version of Python we support.
 # Actually, this is the same version that Python 2.7 supports at the moment!
@@ -102,23 +131,11 @@ def generate_schema(shotgun, schema_file_path, schema_entity_file_path):
     and downloading the schema information for that site. Once the generated schema
     files are being passed to mockgun, it will mimic the site's schema structure.
 
-    :param sg_url: Shotgun site url
-    :param sg_script: Script name to connect with
-    :param sg_key: Script key to connect with
+    :param shotgun: Shotgun instance
     :param schema_file_path: Path where to write the main schema file to
     :param schema_entity_file_path: Path where to write the entity schema file to
     """
 
     schema = shotgun.schema_read()
-    fh = open(schema_file_path, "wb")
-    try:
-        pickle.dump(schema, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL)
-    finally:
-        fh.close()
-
     schema_entity = shotgun.schema_entity_read()
-    fh = open(schema_entity_file_path, "wb")
-    try:
-        pickle.dump(schema_entity, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL)
-    finally:
-        fh.close()
+    SchemaFactory.set_schemas(schema, schema_file_path, schema_entity, schema_entity_file_path)