mlatom 2.0.0__py3-none-any.whl

MLatom/args_class.py

@@ -0,0 +1,396 @@
+#!/usr/bin/env python3
+'''
+  !---------------------------------------------------------------------------! 
+  ! args_class: handling input of MLatom                                      ! 
+  ! Implementations by: Bao-Xin Xue                                           ! 
+  !---------------------------------------------------------------------------! 
+'''
+
+from collections import defaultdict, UserDict
+from os import replace
+from typing import Callable, DefaultDict, Sequence, Tuple, Union
+from typing import List, Tuple, Dict, Set, Any
+from copy import copy
+import re
+import stopper
+
+try:
+    from doc import Doc
+except:
+    Doc = None
+
+
+class KeyNotFoundException(Exception):
+    pass
+
+
+class DuplicatedKeyException(Exception):
+    pass
+
+
+class AttributeDict(UserDict):
+    data: Dict[Any, Any]
+    lower_key_set: Set[Any]
+    lower_key_to_key: Dict[str, str]
+    attribute_dict_variable = ['data', 'lower_key_set', 'lower_key_to_key', 'attribute_dict_variable']
+    
+    def __init__(self) -> None:
+        self.__dict__['data'] = {}
+        self.__dict__['lower_key_set'] = set()
+        self.__dict__['lower_key_to_key'] = {}
+        self.__dict__['writable_flag'] = False
+    
+    @classmethod
+    def merge_dict(cls, my_dict: 'AttributeDict', other: Dict[str, Any]) -> None:
+        for other_key in other.keys():
+            other_key_strip = other_key.strip()        # for the usage of my_dict
+            other_key_lower = other_key_strip.lower()  # for the usage of my_dict.lower_key_set
+            if other_key_lower in my_dict.lower_key_set:
+                if type(other[other_key]) in [dict, cls]:
+                    if type(my_dict[other_key_strip]) is not cls:
+                        my_dict[other_key_strip] = AttributeDict()
+                    cls.merge_dict(my_dict[other_key_strip], other[other_key])
+                else:
+                    my_dict[other_key_strip] = other[other_key]
+            else:
+                my_dict.lower_key_set.add(other_key_lower)
+                my_dict.lower_key_to_key[other_key_lower] = other_key_strip
+                my_dict[other_key_strip] = other[other_key]
+    
+    def key_in_dict(self, name: str) -> bool:
+        lower_key = name.lower().strip()
+        return lower_key in self.lower_key_set
+
+    @classmethod
+    def dict_to_attributedict(cls, d: Dict[str, Any]) -> 'AttributeDict':
+        new_dict = cls()
+        for key, value in d.items():
+            if type(value) is dict:
+                new_value = cls.dict_to_attributedict(value)
+                new_dict[key] = new_value
+            else:
+                new_dict[key] = value
+        return new_dict
+    
+    @classmethod
+    def normal_dict(cls, d: 'AttributeDict') -> Dict[str, Any]:
+        new_dict = {}
+        for key, value in d.items():
+            if type(value) in [cls, dict]:
+                new_dict[key] = cls.normal_dict(value)
+            else:
+                new_dict[key] = value
+        return new_dict
+    
+    def _raise_internal_key_error(self, name: str) -> None:
+        if name.lower() in self.attribute_dict_variable:
+            raise KeyError(name)
+    
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name.startswith('__') or (name in self.__dict__):
+            super(AttributeDict, self).__setattr__(name, value)
+        else:
+            self.__setitem__(name, value)
+    
+    # def __getattribute__(self, name: str) -> Any:
+    #     if name in super(AttributeDict, self).__getattribute__('attribute_dict_variable'):
+    #         raise KeyError(f'{name} can not appears on the args_class.{self.__class__}')
+    #     return super(AttributeDict, self).__getattribute__(name)
+    
+    def __getattr__(self, name: str) -> Any:
+        if not self.key_in_dict(name):
+            self[name] = self.__class__()
+        return self.__getitem__(name)
+    
+    def __setitem__(self, name: str, value: Any) -> Any:
+        self._raise_internal_key_error(name)
+        key_strip = name.strip()
+        key_lower = key_strip.lower()
+        if type(value) is dict:
+            new_value = self.dict_to_attributedict(value)
+        else:
+            new_value = copy(value)
+        if key_lower in self.lower_key_set:
+            true_key = self.lower_key_to_key[key_lower]
+            self.data[true_key] = new_value
+        else:
+            self.lower_key_set.add(key_lower)
+            self.lower_key_to_key[key_lower] = key_strip
+            self.data[key_strip] = new_value
+    
+    def __getitem__(self, name: str) -> Any:
+        self._raise_internal_key_error(name)
+        if name.startswith('__') or (name in self.__dict__):
+            return self.__dict__[name]
+        lower_key = name.lower().strip()
+        if lower_key in self.lower_key_set:
+            return self.data[self.lower_key_to_key[lower_key]]
+        raise KeyError(name)
+
+
+class ArgsBase():
+    """As a argument parsing base class to be extended.
+    
+        basic usage:
+            1. extend this class in your class argument class
+            2. use add_default_dict_args or add_dict_args or add_keyword_value_args to add default key-value pair
+            3. use parse_input_file or parse_input_content to parse argument content
+            4. use self.args_dict to get a argument dict  OR  self.args_string_list(excluse) to get a list of argument string
+        
+        Other feature:
+            1. set_ignore_keyword_list: to exclusive some keyword
+            2. exists_key: to judge whether key exists (for safely check)  (now only support first layer check)
+            3. set_keyword_alias(standard_key, [alias_key1, alias_key2, ...])
+            4. can use add_dict_args to add a dict to the ArgsBase
+        
+        Attention:
+            1. only when set dict_args and read file content can set new argument
+            2. user can use args.xxx = yyy refresh the value, but only when the key exists, the dict will refresh, 
+                OR it will only add a new member var on args, so please combine it with `exists_key` to check it.
+    """
+    ignore_list: List[str] = []
+    keyword_alias_dict: Dict[str, List[str]] = {}
+    writable: bool = True
+    data: AttributeDict
+    def __init__(self) -> None:
+        self.__dict__['data'] = AttributeDict()
+        self.__dict__['writable'] = False
+    
+    def add_default_dict_args(self, keyword_list: List[str], type_or_Callable_or_value: Union[type, Callable[..., Any], Any]):
+        self.writable = True
+        if type(type_or_Callable_or_value) in [type, Callable]:
+            dd: DefaultDict[str, Any] = defaultdict(type_or_Callable_or_value)
+            for key in keyword_list:
+                self.data[key] = dd[key]
+        else:
+            for key in keyword_list:
+                self.data[key] = copy(type_or_Callable_or_value)
+        self.writable = False
+    
+    # @singledispatch
+    def add_dict_args(self, dict: Dict[str, Any]):
+        AttributeDict.merge_dict(self.data, dict)
+    
+    def add_keyword_value_args(self, keyword_list: List[str], value_list: List[Any]):
+        if len(keyword_list) != len(value_list):
+            self._raise_error(f'Internal Error: ArgsBase.add_keyword_value_args: the length is not equal in keyword_list and value_list')
+        d = { k: v for k, v in zip(keyword_list, value_list)}
+        self.add_dict_args(d)
+    
+    @staticmethod
+    def _args_extractor(string: str) -> List[str]:
+        pair_dict = {'(': ')', '[': ']', '{': '}', '<': '>', "'": "'", '"': '"'}
+        pair_level = 0
+        pair_right: List[str] = []
+        tmp = ''
+        prev = ''
+        out_string_list: List[str] = []
+        for chr in string:
+            pair_right_last = pair_right[-1] if pair_right else ''
+            if chr in pair_dict.keys() and chr != pair_right_last:
+                pair_level += 1
+                pair_right.append(pair_dict[chr])
+            elif chr == pair_right_last:
+                pair_level -= 1
+                pair_right.pop()
+            elif chr == ' ' and pair_level == 0 and prev.strip():  # end
+                out_string_list.append(tmp)
+                tmp = ''
+                prev = chr
+                continue
+            if '"' in pair_right or "'" in pair_right:
+                tmp += chr
+            else:
+                tmp += chr.strip()
+            prev = chr
+        if pair_level:
+            stopper.stopMLatom(f'pair character unmatched args: "{tmp}"')
+        else:
+            if tmp: out_string_list.append(tmp)
+        return out_string_list
+    
+    def parse_input_file(self, file: str):
+        try:
+            with open(file) as f:
+                content = f.read().splitlines()
+        except:
+            self._raise_error(f'can not open file {file}! exit!'); exit()
+        self.parse_input_content(content)
+    
+    def parse_input_content(self, content: Union[List[str], str]):
+        help = False
+        if type(content) is str:
+            content = content.splitlines()
+        for c in self._clean_input_content(content):
+            splitted = c.split('=')
+            if len(splitted) == 1:
+                key = splitted[0]
+                if key.lower() in ['help', '-help', '-h', '--help']:
+                    help = True
+                elif key.lower() in self.ignore_list:
+                    pass
+                else:
+                    restored_key = self._restore_alias_key(key)
+                    self.data[restored_key] = True
+            elif len(splitted) == 2:
+                key, value = tuple(map(lambda x: x.strip(), splitted))
+                key, value = self._multi_level_dict(key, value)  # has changed the type of value
+                key = self._restore_alias_key(key)
+                if type(value) in [dict, AttributeDict]:
+                    AttributeDict.merge_dict(self.data, {key: value})
+                    continue
+                else:
+                    self.data[key] = value
+            else:
+                self._raise_error(f'error happends at your input file, error content:\n    {c}')
+        if help:
+            self._print_doc()
+    
+    @classmethod
+    def _multi_level_dict(cls, key: str, value: Any) -> Tuple[str, Union[AttributeDict, Dict[str, Union[str, Dict[str, Any]]]]]:
+        split_key = key.split('.')
+        real_key = split_key[0]
+        inner_value = cls._convert_to_value_with_type(copy(value))
+        for sub_key in split_key[1:][::-1]:
+            outter_dict = AttributeDict()
+            outter_dict[sub_key] = inner_value
+            inner_value = outter_dict
+        return real_key, inner_value
+    
+    def _print_doc(self):
+        try: Doc.printDoc(self.args_dict)
+        except: self._warning_print('Doc not imported')
+        stopper.stopMLatom('')
+    
+    @classmethod
+    def _clean_input_content(cls, content: List[str]) -> List[str]:
+        to_return: List[str] = []
+        for c in content:
+            usable = c.split('#')[0].strip()
+            while usable.find(' =') != -1:
+                usable = usable.replace(' =', '=')
+            while usable.find('= ') != -1:
+                usable = usable.replace('= ', '=')
+            to_return.extend(cls._args_extractor(usable))
+        return to_return
+    
+    @property
+    def args_dict(self) -> Dict[str, Any]:
+        return AttributeDict.normal_dict(self.data)
+    
+    def args_string_list(self, exclusive: List[Any]=[]) -> List[str]:
+        l = []
+        for key, value in self.args_dict.items():
+            if exclusive:
+                if value in exclusive:
+                    continue
+            if value is True:
+                l.append(key)
+            elif value is False:
+                continue
+            else:
+                if type(value) is dict:
+                    l.extend(self._dict_dot_expression(value, key))
+                else:
+                    l.append(f'{key}={value}')
+                # l.append(f'{key}={value}')
+        return l
+    
+    def _dict_dot_expression(self, d: Dict[str, Union[Dict[str, Any], Any]], key_pre_str: str='') -> List[str]:
+        str_list = []
+        for key, value in d.items():
+            if type(value) is dict:
+                str_list.extend(self._dict_dot_expression(value, f'{key_pre_str}.{key}'))
+            else:
+                str_list.append(f'{key_pre_str}.{key}={value}')
+        return str_list
+        
+    
+    def exists_key(self, name: str) -> bool:
+        return self.exists_key(name)
+    
+    def set_ignore_keyword_list(self, kw_or_list: Union[str, List[str]]):
+        if type(kw_or_list) in [list, tuple]:
+            self.ignore_list.extend(kw_or_list)
+        elif type(kw_or_list) is str:
+            self.ignore_list.append(kw_or_list)
+        else:
+            self._raise_error(f'Internal error: "set_ignore_keyword_list" can not be {kw_or_list}')
+    
+    def __getattr__(self, name: str) -> Any:
+        restored_name = self._restore_alias_key(name)
+        return self.data[restored_name]
+    
+    def __setattr__(self, name: str, value: Any) -> None:
+        self.__setitem__(name, value)
+        # restored_name = self._restore_alias_key(name)
+        # if self.data.key_in_dict(restored_name):
+        #     self.data[restored_name] = value
+        # else:
+        #     self.__dict__[name] = value
+    
+    def set_keyword_alias(self, standard: str, alias: Union[str, Sequence[str]]) -> None:
+        if type(alias) is str:
+            self.keyword_alias_dict[standard] = [alias.lower()]
+        else:
+            self.keyword_alias_dict[standard] = list(map(lambda x: x.lower(), alias))
+    
+    def _restore_alias_key(self, key: str) -> str:
+        for std_key, alias_list in self.keyword_alias_dict.items():
+            if key.lower().strip() in alias_list:
+                return std_key
+        return key
+    
+    def __getitem__(self, name: str) -> Any:
+        restored_name = self._restore_alias_key(name)
+        return self.data[restored_name]
+    
+    def __setitem__(self, name: str, value: Any) -> None:
+        restored_name = self._restore_alias_key(name)
+        if self.data.key_in_dict(restored_name):
+            self.data[restored_name] = value
+        else:
+            self.__dict__[name] = value
+    
+    @staticmethod
+    def _raise_error(content: str) -> None:
+        print('WARNING ' * 10)
+        print(' <!> %s <!>' % content)
+        print('WARNING ' * 10)
+        exit()
+    
+    @staticmethod
+    def _warning_print(warning: str) -> None:
+        print('-' * 80)
+        print(' warning:  %s  :warning' % warning)
+        print('-' * 80)
+    
+    @staticmethod
+    def _convert_to_value_with_type(value: Any) -> Any:
+        if type(value) in [dict, AttributeDict]:
+            return value
+        if value.strip().lower() == 'true':
+            return True
+        elif value.strip().lower() == 'false':
+            return False
+        else:
+            try:
+                tmp = int(value)
+            except:
+                try:
+                    tmp = float(value)
+                except:
+                    return value
+                else:
+                    return tmp
+            else:
+                return tmp
+
+
+if __name__ == '__main__':
+    try:
+        args = ArgsBase()
+        args.parse_input_file('inp')
+    except: pass
+    print(  'warning: you are directly running the module: args_class, '\
+            'which should be extended by a class, not run it.')