diff --git a/shotgun_api3/lib/mockgun/__init__.py b/shotgun_api3/lib/mockgun/__init__.py index 6a83e99d..0d891cb6 100644 --- a/shotgun_api3/lib/mockgun/__init__.py +++ b/shotgun_api3/lib/mockgun/__init__.py @@ -31,5 +31,6 @@ """ from .schema import generate_schema # noqa +from .data import generate_data # noqa from .mockgun import Shotgun # noqa from .errors import MockgunError # noqa \ No newline at end of file diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py new file mode 100644 index 00000000..bd0b558a --- /dev/null +++ b/shotgun_api3/lib/mockgun/data.py @@ -0,0 +1,202 @@ +""" + ----------------------------------------------------------------------------- + Copyright (c) 2009-2019, Shotgun Software Inc + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + - Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + + - Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + + - Neither the name of the Shotgun Software Inc nor the names of its + contributors may be used to endorse or promote products derived from this + software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE + FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR + SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER + CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, + OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +----------------------------------------------------------------------------- +""" +# TODO: Python3 support +# TODO: Logging not printing +import cPickle as pickle +import os + +from .errors import MockgunError +from .schema import SchemaFactory + +# 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! +_HIGHEST_24_PICKLE_PROTOCOL = 2 + + +# Global private values to cache the schema in. +__schema, __schema_entity = None, None + + +def _get_schema(force=False): + """ + _get_schema will get the schema from the SchemaFactory and cache it. + + :param bool force: If set, force will always query the latest schema from disk. + :return: schema dictionary from disk + """ + global __schema, __schema_entity + from .mockgun import Shotgun + if not __schema or force is True: + schema_path, schema_entity_path = Shotgun.get_schema_paths() + if not schema_path or not schema_entity_path: + raise MockgunError("You must set the schema paths on the Mockgun instance first.") + __schema, __schema_entity = SchemaFactory.get_schemas(schema_path, schema_entity_path) + return __schema + + +def _get_schema_entity(force=False): + """ + _get_schema_entity will get the schema_entity from the SchemaFactory and cache it. + + :param bool force: If set, force will always query the latest schema_entity from disk. + :return: schema_entity dictionary from disk + """ + global __schema, __schema_entity + from .mockgun import Shotgun + if not __schema_entity or force is True: + schema_path, schema_entity_path = Shotgun.get_schema_paths() + if not schema_path or not schema_entity_path: + raise MockgunError("You must set the schema paths on the Mockgun instance first.") + __schema, __schema_entity = SchemaFactory.get_schemas(schema_path, schema_entity_path) + return __schema_entity + + +def _get_entity_fields(entity): + """ + _get_entity_fields will return a list of the fields on an entity as strings + :param str entity: Shotgun entity that we want the schema for + :return: List of the field names for the provided entity + :rtype: list[str] + """ + schema = _get_schema() + return schema[entity].keys() + + +def _read_data_(shotgun, entity): + """ + _read_data_ will return all of the entries for the provided entity. + It will get all fields for the entity from the Mockgun schema. + + :param shotgun: Shotgun instance used to query a live site + :param str entity: Shotgun entity that we want the schema for + :return: List of found entities + :rtype: list[dict] + """ + try: + return shotgun.find( + entity, + filters=[], + fields=_get_entity_fields(entity) + ) + except Exception as err: + print(" Exception: %s" % str(err)) + import traceback + traceback.print_exc() + return [] + + +class DatabaseFactory(object): + """ + Allows to instantiate a pickled database. + """ + _database_cache = None + _database_cache_path = None + + @classmethod + def get_database(cls, database_path): + """ + Retrieves the schemas from disk. + + :param str database_path: Path to the database. + + :returns: Dictionary holding the database. + :rtype: dict + """ + if not os.path.exists(database_path): + raise MockgunError("Cannot locate Mockgun database file '%s'!" % database_path) + + # Poor man's attempt at a cache. All of our use cases deal with a single pair of files + # for the duration of the unit tests, so keep a cache for both inputs. We don't want + # to deal with ever growing caches anyway. Just having this simple cache has shown + # speed increases of up to 500% for Toolkit unit tests alone. + + if database_path != cls._database_cache_path: + cls._database_cache = cls._read_file(database_path) + cls._database_cache_path = database_path + + return cls._database_cache + + @classmethod + def _read_file(cls, path): + fh = open(path, "rb") + try: + return pickle.load(fh) + 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() + + @classmethod + def set_database(cls, database, database_path): + """ + Writes the schemas to disk. + + :param dict database: The database in memory. + :param str database_path: Path to the database. + """ + if database_path != cls._database_cache_path: + cls._database_cache_path = database_path + cls._database_cache = database + + cls._write_file(database, database_path) + + +# ---------------------------------------------------------------------------- +# Utility methods +def generate_data(shotgun, data_file_path, entity_subset=None): + """ + Helper method for mockgun. + Generates the data files needed by the mocker by connecting to a real shotgun + and downloading the information for that site. Once the generated data + files are being passed to mockgun, it will mimic the site's database structure. + + :param shotgun: Shotgun instance + :param data_file_path: Path where to write the main data file to + :param entity_subset: Optional subset of entities to generate data for. + If not passed, it will default to all entities + """ + + if not entity_subset: + entity_subset = _get_schema().keys() + + database = {} + for entity in entity_subset: + print("Requesting data for: %s" % entity) + database[entity] = _read_data_(shotgun, entity) + + DatabaseFactory.set_database(database, data_file_path) diff --git a/shotgun_api3/lib/mockgun/mockgun.py b/shotgun_api3/lib/mockgun/mockgun.py index d8f33ceb..13bd7d08 100644 --- a/shotgun_api3/lib/mockgun/mockgun.py +++ b/shotgun_api3/lib/mockgun/mockgun.py @@ -120,6 +120,7 @@ from ...shotgun import _Config from .errors import MockgunError from .schema import SchemaFactory +from .data import DatabaseFactory # ---------------------------------------------------------------------------- # Version @@ -147,6 +148,19 @@ class Shotgun(object): __schema_path = None __schema_entity_path = None + __database_path = None + + @classmethod + def set_database_path(cls, database_path): + """ + Set the path where the database can be found. This is done at the class + level so all Shotgun instances will share the same database. + The responsibility to generate and load these files is left to the user + changing the default value. + + :param database_path: Directory path where the database is. + """ + cls.__database_path = database_path @classmethod def set_schema_paths(cls, schema_path, schema_entity_path): @@ -161,6 +175,16 @@ def set_schema_paths(cls, schema_path, schema_entity_path): cls.__schema_path = schema_path cls.__schema_entity_path = schema_entity_path + @classmethod + def get_database_path(cls): + """ + Returns the file which are part of the database. + These path can then be used in generate_database if needed. + + :returns: A string with database_path + """ + return cls.__database_path + @classmethod def get_schema_paths(cls): """ @@ -204,6 +228,7 @@ def __init__(self, # initialize the "database" self._db = dict((entity, {}) for entity in self._schema) + self._update_db() # set some basic public members that exist in the Shotgun API self.base_url = base_url @@ -422,9 +447,20 @@ 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_database(self): + DatabaseFactory.set_database(self._db, self.get_database_path()) + ################################################################################################### # internal methods and members + def _update_db(self): + database = DatabaseFactory.get_database(self.get_database_path()) + for entity_type in database: + for entity in database[entity_type]: + row = self._get_new_row(entity_type) + self._update_row(entity_type, row, entity) + self._db[entity_type][row["id"]] = row + def _validate_entity_type(self, entity_type): if entity_type not in self._schema: raise ShotgunError("%s is not a valid entity" % entity_type) diff --git a/tests/mockgun/database.pickle b/tests/mockgun/database.pickle new file mode 100644 index 00000000..fe195da8 Binary files /dev/null and b/tests/mockgun/database.pickle differ diff --git a/tests/test_mockgun.py b/tests/test_mockgun.py index 2bd4ae10..bc765cd1 100644 --- a/tests/test_mockgun.py +++ b/tests/test_mockgun.py @@ -51,6 +51,9 @@ os.path.join(mockgun_schema_folder, "schema.pickle"), os.path.join(mockgun_schema_folder, "schema_entity.pickle") ) +Mockgun.set_database_path( + os.path.join(mockgun_schema_folder, "database.pickle"), +) # FIXME: This should probably be refactored into a base class for @@ -90,6 +93,7 @@ def test_interface_intact(self): # error. mockgun.MockgunError mockgun.generate_schema + mockgun.generate_data mockgun.Shotgun