__getstate__ and __setstate__ for proto

MyCapytain/resources/proto/inventory.py

@@ -108,6 +108,7 @@ def parse(self, resource):
         """
         raise NotImplementedError()
 
+
 class Text(Resource):
     """ Represents a CTS Text
     """
@@ -148,13 +149,29 @@ def __init__(self, resource=None, urn=None, parents=None, subtype="Edition"):
                 if self.parents[0].texts[urn].subtype == "Edition"
             ]
 
+    def __getstate__(self):
+        """ Pickling method to be called upon dumping object
+
+        :return: Dictionary Representation
+        """
+
+        return dict(
+            metadata=getattr(self.metadata, "__getstate__")(),
+            urn=str(self.urn),
+            lang=self.lang,
+            subtype=self.subtype,
+            parents=[getattr(item, "__getstate__")(children=False) for item in self.parents],
+            citations=[getattr(value, "__getstate__")() for value in self.citation]
+        )
 
 def Edition(resource=None, urn=None, parents=None):
     return Text(resource=resource, urn=urn, parents=parents, subtype="Edition")
 
+
 def Translation(resource=None, urn=None, parents=None):
     return Text(resource=resource, urn=urn, parents=parents, subtype="Translation")
-
+
+
 class Work(Resource):
     """ Represents a CTS Work
     """
@@ -196,6 +213,24 @@ def getLang(self, key=None):
         else:
             return [self.texts[urn] for urn in self.texts if self.texts[urn].subtype == "Translation"]
 
+    def __getstate__(self, children=False):
+        """ Pickling method to be called upon dumping object
+
+        :return: Dictionary Representation
+        """
+        __dict__ = dict(
+            metadata=getattr(self.metadata, "__getstate__")(),
+            urn=str(self.urn),
+            parents=[getattr(item, "__getstate__")(children=False) for item in self.parents]
+        )
+        if children:
+            __dict__["texts"] = {
+                key: getattr(value, "__getstate__")() for key, value in self.texts.items()
+            }
+
+        return __dict__
+
+
 class TextGroup(Resource):
     """ Represents a CTS Textgroup
     """
@@ -223,6 +258,20 @@ def __init__(self, resource=None, urn=None, parents=None):
         if resource is not None:
             self.setResource(resource)
 
+    def __getstate__(self, children=True):
+        """ Pickling method to be called upon dumping object
+
+        :return:
+        """
+        __dict__ = dict(
+            metadata=getattr(self.metadata, "__getstate__")(),
+            urn=str(self.urn)
+        )
+        if children:
+            __dict__["works"] = {
+                key: getattr(value, "__getstate__")() for key, value in self.textgroups.items()
+            }
+        return __dict__
 
 class TextInventory(Resource):
     """ Represents a CTS Inventory file
@@ -239,4 +288,16 @@ def __init__(self, resource=None, id=None):
         self.id = id
         self.parents = ()
         if resource is not None:
-            self.setResource(resource)
+            self.setResource(resource)
+
+    def __getstate__(self):
+        """ Pickling method to be called upon dumping object
+
+        :return:
+        """
+        return {
+            "textgroups": {
+                key: getattr(value, "__getstate__")() for key, value in self.textgroups.items()
+            },
+            "id": self.id
+        }