idds/common/dict_class.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2020 - 2025
0010
0011 """
0012 Dict class.
0013 """
0014
0015 import base64
0016 import json
0017 import logging
0018 import pickle
0019 import urllib
0020 import zlib
0021 import inspect
0022 from enum import Enum
0023
0024
0025 class DictClass(object):
0026     def __init__(self, loading=False):
0027         self._zip_items = []
0028         self._not_auto_unzip_items = []
0029         self._loading = loading
0030
0031     def zip_data(self, data, name=None):
0032         try:
0033             if type(data) in [str] and data.startswith("idds_zip:"):
0034                 # already zipped
0035                 return data
0036
0037             # Convert to JSON string
0038             json_str = json.dumps(data)
0039
0040             # Compress with zlib
0041             compressed = zlib.compress(json_str.encode())
0042
0043             # Encode to base64 for safe storage/transfer
0044             compressed_str = base64.b64encode(compressed).decode()
0045
0046             return "idds_zip:" + compressed_str
0047         except Exception as ex:
0048             print(f"Dict_class failed to zip data: {ex}")
0049         return data
0050
0051     def unzip_data(self, data):
0052         try:
0053             if type(data) not in [str] or not data.startswith("idds_zip:"):
0054                 # not zipped data
0055                 return data
0056
0057             # remove the head 'idds_zip:'
0058             actual_data = data[9:]
0059
0060             # Decode from base64
0061             decoded = base64.b64decode(actual_data)
0062
0063             # Decompress with zlib
0064             decompressed = zlib.decompress(decoded).decode()
0065
0066             # Convert back to dictionary
0067             original_data = json.loads(decompressed)
0068
0069             return original_data
0070         except Exception as ex:
0071             print(f"Dict_class failed to unzip data: {ex}")
0072         return data
0073
0074     @property
0075     def zip_items(self):
0076         if hasattr(self, '_zip_items'):
0077             return self._zip_items
0078         else:
0079             return []
0080
0081     @zip_items.setter
0082     def zip_items(self, value):
0083         self._zip_items = value
0084
0085     @property
0086     def not_auto_unzip_items(self):
0087         if hasattr(self, '_not_auto_unzip_items'):
0088             return self._not_auto_unzip_items
0089         else:
0090             return []
0091
0092     @not_auto_unzip_items.setter
0093     def not_auto_unzip_items(self, value):
0094         self._not_auto_unzip_items = value
0095
0096     def should_zip(self, key):
0097         if key in self.zip_items:
0098             return True
0099         return False
0100
0101     def should_auto_unzip(self, key):
0102         if key in self.zip_items and key not in self.not_auto_unzip_items:
0103             return True
0104         return False
0105
0106     def should_unzip(self, key):
0107         if key in self.zip_items:
0108             return True
0109         return False
0110
0111     def to_dict_l(self, d):
0112         # print('to_dict_l')
0113         # print(d)
0114         if not d:
0115             return d
0116
0117         if hasattr(d, 'refresh'):
0118             d.refresh()
0119
0120         if hasattr(d, 'to_dict'):
0121             return d.to_dict()
0122         elif isinstance(d, dict):
0123             new_d = {}
0124             for k, v in d.items():
0125                 new_d[k] = self.to_dict_l(v)
0126             return new_d
0127         elif isinstance(d, list):
0128             new_d = []
0129             for k in d:
0130                 new_d.append(self.to_dict_l(k))
0131             return new_d
0132         elif inspect.ismethod(d):
0133             return {'idds_method': d.__name__, 'idds_method_class_id': d.__self__.get_internal_id()}
0134         return d
0135
0136     def to_dict(self):
0137         # print('to_dict')
0138         self._loading = False
0139         ret = {'class': self.__class__.__name__,
0140                'module': self.__class__.__module__,
0141                'attributes': {}}
0142
0143         if hasattr(self, 'refresh'):
0144             self.refresh()
0145
0146         for key, value in self.__dict__.items():
0147             # print(key)
0148             # print(value)
0149             # if not key.startswith('__') and not key.startswith('_'):
0150             if not key.startswith('__'):
0151                 if key in ['logger']:
0152                     new_value = None
0153                 elif hasattr(self, 'should_zip') and self.should_zip(key):
0154                     new_value = self.zip_data(value, name=key)
0155                 else:
0156                     new_value = self.to_dict_l(value)
0157                 ret['attributes'][key] = new_value
0158         return ret
0159
0160     @staticmethod
0161     def is_class(d):
0162         if d and isinstance(d, dict) and 'class' in d and 'module' in d and 'attributes' in d:
0163             return True
0164         return False
0165
0166     @staticmethod
0167     def is_class_method(d):
0168         if d and isinstance(d, dict) and 'idds_method' in d and 'idds_method_class_id' in d:
0169             return True
0170         return False
0171
0172     @staticmethod
0173     def is_class_attribute(d):
0174         if d and isinstance(d, dict) and 'idds_attribute' in d and 'idds_method_class_id' in d:
0175             return True
0176         return False
0177
0178     @staticmethod
0179     def load_instance(d):
0180         module = __import__(d['module'], fromlist=[None])
0181         cls = getattr(module, d['class'])
0182         if issubclass(cls, Enum):
0183             impl = cls(d['attributes']['_value_'])
0184         else:
0185             sig = inspect.signature(cls.__init__)
0186             if 'json_load' in sig.parameters and 'loading' in sig.parameters:
0187                 impl = cls(json_load=True, loading=True)
0188             elif 'json_load' in sig.parameters:
0189                 impl = cls(json_load=True)
0190             elif 'loading' in sig.parameters:
0191                 impl = cls(loading=True)
0192             else:
0193                 impl = cls()
0194             impl._loading = False
0195         return impl
0196
0197     @staticmethod
0198     def load_instance_method(d):
0199         # not do anything. Will load the method in Workflow class.
0200         return d
0201
0202     @staticmethod
0203     def load_instance_attribute(d):
0204         # not do anything. Will load the method in Workflow class.
0205         return d
0206
0207     @staticmethod
0208     def from_dict(d):
0209         if not d:
0210             return d
0211
0212         # print("from_dict: %s" % str(d))
0213         # print("is_class: %s" % DictClass.is_class(d))
0214         if isinstance(d, DictBase):
0215             d.metadata = d.metadata
0216
0217         if DictClass.is_class(d):
0218             impl = DictClass.load_instance(d)
0219             impl._loading = True
0220             last_items = {}
0221             for key, value in d['attributes'].items():
0222                 # print(key)
0223                 if key in ['logger']:
0224                     continue
0225                 elif key == "_metadata":
0226                     last_items[key] = value
0227                 # elif key == 'output_data':
0228                 #     continue
0229                 else:
0230                     value = DictClass.from_dict(value)
0231                 setattr(impl, key, value)
0232
0233             # unzip
0234             for key, value in impl.__dict__.items():
0235                 if hasattr(impl, 'should_auto_unzip') and impl.should_auto_unzip(key):
0236                     new_value = impl.unzip_data(value)
0237                     setattr(impl, key, new_value)
0238
0239             # print("last_items: %s" % str(last_items))
0240             for key, value in last_items.items():
0241                 value = DictClass.from_dict(value)
0242                 setattr(impl, key, value)
0243
0244             impl._loading = False
0245             return impl
0246         elif DictClass.is_class_method(d):
0247             impl = DictClass.load_instance_method(d)
0248             return impl
0249         elif DictClass.is_class_attribute(d):
0250             impl = DictClass.load_instance_attribute(d)
0251             return impl
0252         elif isinstance(d, dict):
0253             for k, v in d.items():
0254                 d[k] = DictClass.from_dict(v)
0255             return d
0256         elif isinstance(d, list):
0257             new_d = []
0258             for k in d:
0259                 new_d.append(DictClass.from_dict(k))
0260             return new_d
0261         else:
0262             return d
0263
0264         return d
0265
0266
0267 class DictMetadata(DictClass):
0268     def __init__(self, loading=False):
0269         super(DictMetadata, self).__init__(loading=loading)
0270         pass
0271
0272     def add_item(self, key, value):
0273         setattr(self, key, value)
0274
0275     def get_item(self, key, default):
0276         return getattr(self, key, default)
0277
0278
0279 class DictBase(DictClass):
0280     def __init__(self, loading=False):
0281         super(DictBase, self).__init__(loading=loading)
0282         self.metadata = DictMetadata()
0283         pass
0284
0285     def add_metadata_item(self, key, value):
0286         self.metadata.add_item(key, value)
0287
0288     def get_metadata_item(self, key, default=None):
0289         return self.metadata.get_item(key, default)
0290
0291     def refresh(self):
0292         pass
0293
0294     def load_metadata(self):
0295         pass
0296
0297     @property
0298     def metadata(self):
0299         return self._metadata
0300
0301     @metadata.setter
0302     def metadata(self, value):
0303         self._metadata = value
0304         self.load_metadata()
0305
0306     def IDDSProperty(self, attribute):
0307         def _get(self, attribute):
0308             self.get_metadata_item(attribute, None)
0309
0310         def _set(self, attribute, value):
0311             self.add_metadata_item(attribute, value)
0312
0313         attribute = property(_get, _set)
0314         return attribute
0315
0316     def serialize(self):
0317         return urllib.parse.quote_from_bytes(pickle.dumps(self))
0318
0319     @staticmethod
0320     def deserialize(obj):
0321         # return urllib.parse.unquote_to_bytes(pickle.loads(obj))
0322         return pickle.loads(urllib.parse.unquote_to_bytes(obj))
0323
0324     def get_class_name(self):
0325         return self.__class__.__name__
0326
0327     def setup_logger(self):
0328         """
0329         Setup logger
0330         """
0331         self.logger = logging.getLogger(self.get_class_name())