pandaharvester/harvestercore/spec_base.py

0001 """
0002 Base class for XyzSpec
0003
0004 """
0005
0006 import json
0007 import pickle
0008 from json.decoder import JSONDecodeError
0009
0010 import rpyc
0011
0012
0013 # encoder for non-native json objects
0014 class PythonObjectEncoder(json.JSONEncoder):
0015     def default(self, obj):
0016         if isinstance(obj, rpyc.core.netref.BaseNetref):
0017             retVal = rpyc.utils.classic.obtain(obj)
0018         else:
0019             retVal = {"_non_json_object": pickle.dumps(obj)}
0020         return retVal
0021
0022
0023 # hook for decoder
0024 def as_python_object(dct):
0025     if "_non_json_object" in dct:
0026         return pickle.loads(str(dct["_non_json_object"]))
0027     return dct
0028
0029
0030 # base class for XyzSpec
0031 class SpecBase(object):
0032     # to be set
0033     attributesWithTypes = ()
0034     zeroAttrs = ()
0035     skipAttrsToSlim = ()
0036
0037     # constructor
0038     def __init__(self):
0039         # remove types
0040         object.__setattr__(self, "attributes", [])
0041         object.__setattr__(self, "serializedAttrs", set())
0042         for attr in self.attributesWithTypes:
0043             attr, attrType = attr.split(":")
0044             attrType = attrType.split()[0]
0045             self.attributes.append(attr)
0046             if attrType in ["blob"]:
0047                 self.serializedAttrs.add(attr)
0048         # install attributes
0049         for attr in self.attributes:
0050             if attr in self.zeroAttrs:
0051                 object.__setattr__(self, attr, 0)
0052             else:
0053                 object.__setattr__(self, attr, None)
0054         # map of changed attributes
0055         object.__setattr__(self, "changedAttrs", {})
0056
0057     # override __setattr__ to collect changed attributes
0058     def __setattr__(self, name, value):
0059         oldVal = getattr(self, name)
0060         object.__setattr__(self, name, value)
0061         newVal = getattr(self, name)
0062         # collect changed attributes
0063         if oldVal != newVal:
0064             self.changedAttrs[name] = value
0065
0066     # keep state for pickle
0067     def __getstate__(self):
0068         odict = self.__dict__.copy()
0069         del odict["changedAttrs"]
0070         return odict
0071
0072     # restore state from the unpickled state values
0073     def __setstate__(self, state):
0074         self.__init__()
0075         for k, v in state.items():
0076             object.__setattr__(self, k, v)
0077
0078     # reset changed attribute list
0079     def reset_changed_list(self):
0080         object.__setattr__(self, "changedAttrs", {})
0081
0082     # force update
0083     def force_update(self, name):
0084         if name in self.attributes:
0085             self.changedAttrs[name] = getattr(self, name)
0086
0087     # force not update
0088     def force_not_update(self, name):
0089         if name in self.changedAttrs:
0090             del self.changedAttrs[name]
0091
0092     # check if attributes are updated
0093     def has_updated_attributes(self):
0094         return len(self.changedAttrs) > 0
0095
0096     # pack into attributes
0097     def pack(self, values, slim=False):
0098         if hasattr(values, "_asdict"):
0099             values = values._asdict()
0100         for attr in self.attributes:
0101             if slim and attr in self.skipAttrsToSlim:
0102                 val = None
0103             else:
0104                 val = values[attr]
0105                 if attr in self.serializedAttrs and val is not None:
0106                     try:
0107                         val = json.loads(val, object_hook=as_python_object)
0108                     except JSONDecodeError:
0109                         pass
0110             object.__setattr__(self, attr, val)
0111
0112     # set blob attribute
0113     def set_blob_attribute(self, key, val):
0114         try:
0115             val = json.loads(val, object_hook=as_python_object)
0116             object.__setattr__(self, key, val)
0117         except JSONDecodeError:
0118             pass
0119
0120     # return column names for INSERT
0121     def column_names(cls, prefix=None, slim=False):
0122         ret = ""
0123         for attr in cls.attributesWithTypes:
0124             attr = attr.split(":")[0]
0125             if slim and attr in cls.skipAttrsToSlim:
0126                 continue
0127             if prefix is None:
0128                 ret += f"{attr},"
0129             else:
0130                 ret += f"{prefix}.{attr},"
0131         ret = ret[:-1]
0132         return ret
0133
0134     column_names = classmethod(column_names)
0135
0136     # return expression of bind variables for INSERT
0137     def bind_values_expression(cls):
0138         ret = "VALUES("
0139         for attr in cls.attributesWithTypes:
0140             attr = attr.split(":")[0]
0141             ret += f":{attr},"
0142         ret = ret[:-1]
0143         ret += ")"
0144         return ret
0145
0146     bind_values_expression = classmethod(bind_values_expression)
0147
0148     # return an expression of bind variables for UPDATE to update only changed attributes
0149     def bind_update_changes_expression(self):
0150         ret = ""
0151         for attr in self.attributes:
0152             if attr in self.changedAttrs:
0153                 ret += f"{attr}=:{attr},"
0154         ret = ret[:-1]
0155         ret += " "
0156         return ret
0157
0158     # return map of values
0159     def values_map(self, only_changed=False):
0160         ret = {}
0161         for attr in self.attributes:
0162             # only changed attributes
0163             if only_changed:
0164                 if attr not in self.changedAttrs:
0165                     continue
0166             val = getattr(self, attr)
0167             if val is None:
0168                 if attr in self.zeroAttrs:
0169                     val = 0
0170                 else:
0171                     val = None
0172             if attr in self.serializedAttrs:
0173                 val = json.dumps(val, cls=PythonObjectEncoder)
0174             ret[f":{attr}"] = val
0175         return ret
0176
0177     # return list of values
0178     def values_list(self, only_changed=False):
0179         ret = []
0180         for attr in self.attributes:
0181             # only changed attributes
0182             if only_changed:
0183                 if attr not in self.changedAttrs:
0184                     continue
0185             val = getattr(self, attr)
0186             if val is None:
0187                 if attr in self.zeroAttrs:
0188                     val = 0
0189                 else:
0190                     val = None
0191             if attr in self.serializedAttrs:
0192                 val = json.dumps(val, cls=PythonObjectEncoder)
0193             ret.append(val)
0194         return ret
0195
0196     # get dict of changed attributes
0197     def get_changed_attributes(self):
0198         retDict = dict()
0199         for attr in self.changedAttrs:
0200             retDict[attr] = getattr(self, attr)
0201         return retDict
0202
0203     # set attributes
0204     def set_attributes_with_dict(self, attr_dict):
0205         for attr, val in attr_dict.items():
0206             setattr(self, attr, val)