From 15fa1dd485f713837b50d4b50ebae9ed3a3f2a0f Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 12:27:45 -0700 Subject: [PATCH 01/13] Add database_path concept to Mockgun The idea is that we will be able to serialize our database into a mockgun-able format. We have a "database_path" like the schema_paths. --- shotgun_api3/lib/mockgun/mockgun.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/shotgun_api3/lib/mockgun/mockgun.py b/shotgun_api3/lib/mockgun/mockgun.py index d8f33cebb..16d5f7acd 100644 --- a/shotgun_api3/lib/mockgun/mockgun.py +++ b/shotgun_api3/lib/mockgun/mockgun.py @@ -147,6 +147,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 +174,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): """ From f2c468201dcc7140918f27819d5e557073f5a358 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 12:28:27 -0700 Subject: [PATCH 02/13] Load the testing database for Mockgun Just has a bunch of the default projects from a shotgun site setup. --- tests/test_mockgun.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_mockgun.py b/tests/test_mockgun.py index 2bd4ae108..9cf298c32 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, "data.pickle"), +) # FIXME: This should probably be refactored into a base class for From d3430e52457e9c89513e0a1c0816e5929c33c404 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 12:29:55 -0700 Subject: [PATCH 03/13] Add generate_data to mockgun interface test --- tests/test_mockgun.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_mockgun.py b/tests/test_mockgun.py index 9cf298c32..c9c592ace 100644 --- a/tests/test_mockgun.py +++ b/tests/test_mockgun.py @@ -93,6 +93,7 @@ def test_interface_intact(self): # error. mockgun.MockgunError mockgun.generate_schema + mockgun.generate_data mockgun.Shotgun From 14666837d086bd9d49447c15b509a2694bb3e290 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 12:59:05 -0700 Subject: [PATCH 04/13] Update database pickle name for test Also git add the database pickle --- tests/mockgun/database.pickle | Bin 0 -> 14481 bytes tests/test_mockgun.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 tests/mockgun/database.pickle diff --git a/tests/mockgun/database.pickle b/tests/mockgun/database.pickle new file mode 100644 index 0000000000000000000000000000000000000000..fe195da87c14d4fb30450d771e0e40809e928e07 GIT binary patch literal 14481 zcmeHOZFk#57KS8k5{I^F!&~W=br;GOHg;x4FEc6|x~6GKfd;mTXklT!qtPf4wJqzd zrJUn)_R}8v6Zyn)_#=CzSWaXsiPhrf>^b?6M3F6N=HC0<``j1(d#dNZY|efec;8DA zedkYI>-p2y%=w6S9yU9?EjP?rH*DIn?L9S539sb^<^m?ha{tGi54W3P#G^ReFe`D_ z=8?3UJQB%4PUPo3w z_XKZkW&io7-oNLc{?9#s)||8BfJd&^*)V5>XT$2|Ew|OOJRaB^=91uI2Oe*`VYkJf zr0X{qw`3dsi)6LGN%L&P!-vhNe>(imJQasBfWN1^J3N$d@Rn$~vJ*8!Hdu z4x1r=ByEs_+veNw&!4+y&O}eT(wtw9+q`4K4b3Un{_IV2CXu3fZgY3L z)y=ru^@1o)KfPhTCgLDS9+fPi>F`Hh0MC4CUT*O)YH}fD7{ZRMC(Zp|gx#kP%k*Ye zKm|e6Ok}-bo&}ZIG8BQ^O+;nVtsOV4c6U`=g+q9jU5&hI7Y7*W zcoCdw$w*!?&qo0l4`J^E8T9;lbH2kLi5>2C!k%BbX3pM>qdOo>^FsfiF5;*TTXLf( z<}{o$EAT5jF_Xx_Tugq0!r{JtEIUG)ry8)&$XrVQB%ko)|7NnwINUR=MAQqVMLj1w zU=*orW+J`~3ZCPE*l|IfKlSx9Y9#~p@@OZq+@Gfxrj&QUa(5C_e8pT!&7c{2aUf)K zcfT9vImd0a!w5v*?3>8~z$de{pUlM;@7Re%yI@1jIB1#gt|W&mef2cGM>2Ti%Evdd z6Ry3gE9P0SXD^PL{j-2dkWnI?jkxEZISQ~uQqVe$qYyodIrN+_iJoC3TqwQ$Ly+*akAwpiI1Gqi zW-c$BUC98k^g;n)91I9^_BOa>Fg*aqtHr?hDT`zk4J~8U0Bk7xaABTqhzOj+p4;iW zfL~{(7havtcKq6*@aXxkgLYEif5Eb><{(Vr`uE*!a~hz(b#V&3NfdR%>nkhacC9UK zm)GLZ3u*}0L#39u|CWS+C+ob;Au{sEVO@Cb6@@trA%`-_h=VXwh@vXCv?;x8lGa1LIsx8e5Eh<)XkffG@Pt_-<8|_;9qG(5%8Thh#?^pceh?P=UXn& zwv-y2j=L?7+i@`|I=2Rd?s~P+fxD(0;dkLgTT2r+Xu~lqtPA2`fi;y|+)*fInCn7t zMR%Ag1w*#5Q-h9n;^zupnd!Jg7%xapnb7q&^Di&@@|-8gzXVTkA&-m zsf7;);rgw@aNYASWy%?(7Q4dlXE8%=BR+!fK7BU6OVMz6@A-cyT-c3+=zBaM$ou)g z8qJS+rz3-Ux6@|==O^ws9cmLDID_JFMjZ_? zL#50>RNAU-5tT5eS{g>iN`84|{AT{j6SpZ0;wYM?sD@Bp<+lu%t@Vkh)O zRW*Zho$8Y6K*%`Lgl6fMLKzui0Q0*HAhDU0q&7q7d+zu4~*P zy6xC1GZYj2D#n$x z_@v-vX6%csxq@-O{`B^lcLoXe)k8z-0p;E^$&w~pb$TyTNvT@z#_^f`pL1q^?YYdp z=f7Wi`(x18A+-%-<-gOFMM5*Id{Bs$2^C){8wNnn|4Sx<97`?)WmcyhU}gCTtn7AI zSW>XnjsQ!)m68@~YtICh+fc2ZELiH9s7it5-tmEDB?lH%5-iW3?e%ycQ;!6iMwTJ-0ifX{fCf=%DU*%% zEI_-k^;LCK2CfVzi#9!zTPd_{9v|9h4sGn&&~|0Gow&T&?<>}2*A1ahd{}?coXfs# zE@eka2eN-@p7r&7qqK;}KxdJC6hd@hhSy&xL`NEsJ^#Z@4rL)acPrJ2D}XruwHTv= z2F+YT91)ac_5pOv_RQ*`IPLi#!TQto%$qx58ye9GIQP0Vdrg|F6EDq3$YG)A&!Cs1 z*fw}0Vla*^hM10oN{LRHqdCycIcT9}iZln^<0R8Xj(}2vW>_|621isFhGHen8z?(4 zh1x1~*^qE-=$Pry+kwPaaj=vEV;iat(4#Ngu=%A#z2@Rj<+b`4x?@!stO$Rn^3CKc zkYL3+S5+YibhnVopk$ycmg7JJlR`&P>3ggDX!=8c7|>G*Y9` zBU2D{2$z&N2Gt}L7FD2NLmjSx9uQWDWkauqA#OmKSd~y&Q*pj=b0hQkrO1VOeB|P9 zIdai>HgfSCvr1!?hz>#a2m_cqw{WSJ5z&te89>j!nW-rkn9Y8)JNR)EDj+d!l-{ou zF0mZFzg0{f_G(U4T_#w8nJ3_w$*9@rkYZYM`fd8mLj7_k)Z2%$h68F&7t@;47vPMO zVKA8PI-2DBpY;6oObX-Ti9gST;+?WjON+H%9Upo>&7t?RlIYFT_VJK-uo{Pv2u4%6 zbR{#I&kK<_p?Vb4NAYww(`^p&FeRE2+8Cp1OEgqjZHY7`;eQLrN&&<4yxX)3k3W)YbZa*J9- zP(=kg!x+$RjY-De#~#y3UXu8xnB-*n0=!9dXp#Vo=6WW!e;mqa4oDK- zN<#l`1gWm=n(LA8M)=?M{O__IPbQ)ay~p&=vQJAL!Gjm!2#VjRXy%aouVX+mFDVeY&>R<;JjQh%S&MuF#_{X_15)vW)c^nh literal 0 HcmV?d00001 diff --git a/tests/test_mockgun.py b/tests/test_mockgun.py index c9c592ace..bc765cd17 100644 --- a/tests/test_mockgun.py +++ b/tests/test_mockgun.py @@ -52,7 +52,7 @@ os.path.join(mockgun_schema_folder, "schema_entity.pickle") ) Mockgun.set_database_path( - os.path.join(mockgun_schema_folder, "data.pickle"), + os.path.join(mockgun_schema_folder, "database.pickle"), ) From e3bedebfffda08d82d78951d44d82cbcb332fa51 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:17:31 -0700 Subject: [PATCH 05/13] Expose generate_data to mockgun interface --- shotgun_api3/lib/mockgun/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/shotgun_api3/lib/mockgun/__init__.py b/shotgun_api3/lib/mockgun/__init__.py index 6a83e99d7..0d891cb6b 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 From d4855e43e1898eed06af61082baa553dc36e7fd1 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:18:04 -0700 Subject: [PATCH 06/13] Include DatabaseFactory and allow database loading to Mockgun --- shotgun_api3/lib/mockgun/mockgun.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/shotgun_api3/lib/mockgun/mockgun.py b/shotgun_api3/lib/mockgun/mockgun.py index 16d5f7acd..b3bfa1983 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 @@ -227,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 @@ -448,6 +450,14 @@ def upload_thumbnail(self, entity_type, entity_id, path, **kwargs): ################################################################################################### # 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) From 2649a6fcdaf8b0b54aa4b4b04ee17704044d84be Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:21:46 -0700 Subject: [PATCH 07/13] Create data module for generating live data duplication DatabaseFactory allows us to load the database from disk into memory generate_data works the same as `generate_schema` but allows an `entity_subset` --- shotgun_api3/lib/mockgun/data.py | 150 +++++++++++++++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 shotgun_api3/lib/mockgun/data.py diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py new file mode 100644 index 000000000..1bc74ee1c --- /dev/null +++ b/shotgun_api3/lib/mockgun/data.py @@ -0,0 +1,150 @@ +""" + ----------------------------------------------------------------------------- + Copyright (c) 2009-2017, 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. + +----------------------------------------------------------------------------- +""" + +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 + + +__schema, __schema_entity = None, None + + +def _get_schema(): + global __schema, __schema_entity + from .mockgun import Shotgun + if not __schema: + __schema, __schema_entity = SchemaFactory.get_schemas(*Shotgun.get_schema_paths()) + return __schema + + +def _get_schema_entity(): + global __schema, __schema_entity + from .mockgun import Shotgun + if not __schema_entity: + __schema, __schema_entity = SchemaFactory.get_schemas(*Shotgun.get_schema_paths()) + return __schema_entity + + +def _get_entity_fields(entity): + schema = _get_schema() + return schema[entity].keys() + + +def _read_data_(shotgun, entity): + 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() + + +# ---------------------------------------------------------------------------- +# 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) + + fh = open(data_file_path, "wb") + try: + pickle.dump(database, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL) + finally: + fh.close() + From 69f31ad97d1678232b8ef147f440642dd08ace96 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:22:47 -0700 Subject: [PATCH 08/13] Add TODO messages --- shotgun_api3/lib/mockgun/data.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 1bc74ee1c..52f4ebb95 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -29,7 +29,9 @@ ----------------------------------------------------------------------------- """ - +# TODO: Python3 support +# TODO: Logging not printing +# TODO: Dump database ability import cPickle as pickle import os From 3e4ee098f00397e9f1fcc8271d970e55709e7450 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:31:59 -0700 Subject: [PATCH 09/13] Allow force updating schema Also bugfixes confirming that the client has set the schema paths first. --- shotgun_api3/lib/mockgun/data.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 52f4ebb95..3db64a08a 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -49,16 +49,22 @@ def _get_schema(): global __schema, __schema_entity from .mockgun import Shotgun - if not __schema: - __schema, __schema_entity = SchemaFactory.get_schemas(*Shotgun.get_schema_paths()) + 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(): global __schema, __schema_entity from .mockgun import Shotgun - if not __schema_entity: - __schema, __schema_entity = SchemaFactory.get_schemas(*Shotgun.get_schema_paths()) + 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 From 868baee946d38b05aa7d1f18bc89af157d678e62 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:32:14 -0700 Subject: [PATCH 10/13] Add docstrings --- shotgun_api3/lib/mockgun/data.py | 32 ++++++++++++++++++++++++++++++-- 1 file changed, 30 insertions(+), 2 deletions(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 3db64a08a..8aec1ce6c 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -43,10 +43,17 @@ _HIGHEST_24_PICKLE_PROTOCOL = 2 +# Global private values to cache the schema in. __schema, __schema_entity = None, None -def _get_schema(): +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: @@ -57,7 +64,13 @@ def _get_schema(): return __schema -def _get_schema_entity(): +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: @@ -69,11 +82,26 @@ def _get_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, From b055b48809a0b688dcd1a50c36bb501909aac339 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:32:50 -0700 Subject: [PATCH 11/13] Update license --- shotgun_api3/lib/mockgun/data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 8aec1ce6c..8c15caaa5 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -1,6 +1,6 @@ """ ----------------------------------------------------------------------------- - Copyright (c) 2009-2017, Shotgun Software Inc + 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: From 78864ffaaa1fff55c884f81439a400ecbfaef586 Mon Sep 17 00:00:00 2001 From: Alex Date: Fri, 3 May 2019 13:39:39 -0700 Subject: [PATCH 12/13] Implement database dumping The functionality now lives on the DatabaseFactory as `set_database`. --- shotgun_api3/lib/mockgun/data.py | 30 ++++++++++++++++++++++------- shotgun_api3/lib/mockgun/mockgun.py | 3 +++ 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 8c15caaa5..71b38baf3 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -31,7 +31,6 @@ """ # TODO: Python3 support # TODO: Logging not printing -# TODO: Dump database ability import cPickle as pickle import os @@ -154,6 +153,28 @@ def _read_file(cls, path): finally: fh.close() + @classmethod + def _write_file(cls, data, path): + fh = open(path, "rb") + 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 @@ -178,9 +199,4 @@ def generate_data(shotgun, data_file_path, entity_subset=None): print("Requesting data for: %s" % entity) database[entity] = _read_data_(shotgun, entity) - fh = open(data_file_path, "wb") - try: - pickle.dump(database, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL) - finally: - fh.close() - + DatabaseFactory.set_database(database, data_file_path) diff --git a/shotgun_api3/lib/mockgun/mockgun.py b/shotgun_api3/lib/mockgun/mockgun.py index b3bfa1983..13bd7d08e 100644 --- a/shotgun_api3/lib/mockgun/mockgun.py +++ b/shotgun_api3/lib/mockgun/mockgun.py @@ -447,6 +447,9 @@ 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 From 02e11410f1ba7c576261021cc38e993dd954f95e Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 21 May 2019 20:51:33 +0100 Subject: [PATCH 13/13] Fix wrong file descriptor --- shotgun_api3/lib/mockgun/data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/shotgun_api3/lib/mockgun/data.py b/shotgun_api3/lib/mockgun/data.py index 71b38baf3..bd0b558a0 100644 --- a/shotgun_api3/lib/mockgun/data.py +++ b/shotgun_api3/lib/mockgun/data.py @@ -155,7 +155,7 @@ def _read_file(cls, path): @classmethod def _write_file(cls, data, path): - fh = open(path, "rb") + fh = open(path, "wb") try: return pickle.dump(data, fh, protocol=_HIGHEST_24_PICKLE_PROTOCOL) finally: