konfai 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

konfai/utils/config.py

@@ -1,17 +1,20 @@
-import os
-import ruamel.yaml
-import inspect 
 import collections
+import inspect
+import os
+import types
 from copy import deepcopy
-from typing import Union, Literal, get_origin, get_args, Any
+from typing import Any, Literal, Union, get_args, get_origin
+
+import ruamel.yaml
 import torch
-from konfai import CONFIG_FILE
+
+from konfai import config_file
 from konfai.utils.utils import ConfigError
 
 yaml = ruamel.yaml.YAML()
 
 
-class Config():
+class Config:
 
     def __init__(self, filename, key) -> None:
         self.filename = filename
@@ -24,32 +27,30 @@ class Config():
                 os.environ["KONFAI_CONFIG_MODE"] = "interactive" if result == "interactive" else "default"
             else:
                 exit(0)
-            with open(self.filename, "w") as f:
-                pass
-                    
-        self.yml = open(self.filename, 'r')
+
+        self.yml = open(self.filename)
         self.data = yaml.load(self.yml)
-        if self.data == None:
+        if self.data is None:
             self.data = {}
-        
+
         self.config = self.data
 
         for key in self.keys:
-            if self.config == None or key not in self.config:
-                self.config = {key : {}}
-            
+            if self.config is None or key not in self.config:
+                self.config = {key: {}}
+
             self.config = self.config[key]
         return self
 
-    def createDictionary(self, data, keys, i) -> dict:
+    def create_dictionary(self, data, keys, i) -> dict:
         if keys[i] not in data:
             data = {keys[i]: data}
         if i == 0:
             return data
         else:
             i -= 1
-            return self.createDictionary(data, keys, i)
-    
+            return self.create_dictionary(data, keys, i)
+
     def merge(self, dict1, dict2) -> dict:
         result = deepcopy(dict1)
 
@@ -57,29 +58,34 @@ class Config():
             if isinstance(value, collections.abc.Mapping):
                 result[key] = self.merge(result.get(key, {}), value)
             else:
-                if not dict2[key] == None:
+                if dict2[key] is not None:
                     result[key] = deepcopy(dict2[key])
         return result
 
-
-    def __exit__(self, type, value, traceback) -> None:
+    def __exit__(self, exc_type, value, traceback) -> None:
         self.yml.close()
         if os.environ["KONFAI_CONFIG_MODE"] == "remove":
-            if os.path.exists(CONFIG_FILE()):
-                os.remove(CONFIG_FILE())
+            if os.path.exists(config_file()):
+                os.remove(config_file())
             return
-        with open(self.filename, 'r') as yml:
+        with open(self.filename) as yml:
             data = yaml.load(yml)
-            if data == None:
+            if data is None:
                 data = {}
-        with open(self.filename, 'w') as yml:
-            yaml.dump(self.merge(data, self.createDictionary(self.config, self.keys, len(self.keys)-1)), yml)
-  
+        with open(self.filename, "w") as yml:
+            yaml.dump(
+                self.merge(
+                    data,
+                    self.create_dictionary(self.config, self.keys, len(self.keys) - 1),
+                ),
+                yml,
+            )
+
     @staticmethod
-    def _getInput(name : str, default : str) ->  str:
+    def _get_input(name: str, default: str) -> str:
         try:
-            return input("{} [{}]: ".format(name, ",".join(default.split(":")[1:]) if len(default.split(":")) else ""))
-        except:
+            return input(f"{name} [{','.join(default.split(':')[1:]) if ':' in default else ''}]: ")
+        except Exception:
             result = input("\nKeep a default configuration file ? (yes,no) : ")
             if result == "yes":
                 os.environ["KONFAI_CONFIG_MODE"] = "default"
@@ -89,119 +95,143 @@ class Config():
         return default.split(":")[1] if len(default.split(":")) > 1 else default
 
     @staticmethod
-    def _getInputDefault(name : str, default : Union[str, None], isList : bool = False) -> Union[list[Union[str, None]], str, None]:
-        if isinstance(default, str) and (default == "default" or (len(default.split(":")) > 1 and default.split(":")[0] == "default")):
+    def _get_input_default(name: str, default: str | None, is_list: bool = False) -> list[str | None] | str | None:
+        if isinstance(default, str) and (
+            default == "default" or (len(default.split(":")) > 1 and default.split(":")[0] == "default")
+        ):
             if os.environ["KONFAI_CONFIG_MODE"] == "interactive":
-                if isList:
-                    list_tmp = []
+                if is_list:
+                    list_tmp: list[str | None] = []
                     key_tmp = "OK"
                     while (key_tmp != "!" and key_tmp != " ") and os.environ["KONFAI_CONFIG_MODE"] == "interactive":
-                        key_tmp = Config._getInput(name, default)
-                        if (key_tmp != "!" and key_tmp != " "):
+                        key_tmp = Config._get_input(name, default)
+                        if key_tmp != "!" and key_tmp != " ":
                             if key_tmp == "":
                                 key_tmp = default.split(":")[1] if len(default.split(":")) > 1 else default
                             list_tmp.append(key_tmp)
                     return list_tmp
                 else:
-                    value = Config._getInput(name, default)
+                    value = Config._get_input(name, default)
                     if value == "":
-                        return default.split(":")[1] if len(default.split(":")) > 1 else default 
-                    else: 
+                        return default.split(":")[1] if len(default.split(":")) > 1 else default
+                    else:
                         return value
             else:
                 default = default.split(":")[1] if len(default.split(":")) > 1 else default
-        return  [default] if isList else default
+        return [default] if is_list else default
 
-    def getValue(self, name, default) -> object:  
+    def get_value(self, name, default) -> object:
         if name in self.config and self.config[name] is not None:
             value = self.config[name]
-            if value == None:
+            if value is None:
                 value = default
             value_config = value
         else:
-            value = Config._getInputDefault(name, default if default != inspect._empty else None)
-            
+            value = Config._get_input_default(name, default if default != inspect._empty else None)
+
             value_config = value
-            if type(value_config) == tuple:
+            if isinstance(value_config, tuple):
                 value_config = list(value)
-                
-            if type(value_config) == list:
-                list_tmp = [] 
+
+            if isinstance(value_config, list):
+                list_tmp = []
                 for key in value_config:
-                    list_tmp.extend(Config._getInputDefault(name, key, isList=True))
-                
+                    res = Config._get_input_default(name, key, is_list=True)
+                    if isinstance(res, list):
+                        list_tmp.extend(res)
+                    else:
+                        list_tmp.append(str(res))
 
                 value = list_tmp
                 value_config = list_tmp
 
-            if type(value) == dict:
+            if isinstance(value, dict):
                 key_tmp = []
-                
+
                 value_config = {}
                 dict_value = {}
                 for key in value:
-                    key_tmp.extend(Config._getInputDefault(name, key, isList=True))
+                    res = Config._get_input_default(name, key, is_list=True)
+                    if isinstance(res, list):
+                        key_tmp.extend(res)
+                    else:
+                        key_tmp.append(str(res))
                 for key in key_tmp:
                     if key in value:
                         value_tmp = value[key]
                     else:
-                        value_tmp = next(v for k,v in value.items() if "default" in k)
+                        value_tmp = next(v for k, v in value.items() if "default" in k)
 
                     value_config[key] = None
                     dict_value[key] = value_tmp
                 value = dict_value
         if isinstance(self.config, str):
-            os.environ['KONFAI_CONFIG_VARIABLE'] = "True"
+            os.environ["KONFAI_CONFIG_VARIABLE"] = "True"
             return None
 
         self.config[name] = value_config if value_config is not None else "None"
         if value == "None":
             value = None
         return value
-                    
-def config(key : Union[str, None] = None):
+
+
+def config(key: str | None = None):
     def decorator(function):
         def new_function(*args, **kwargs):
             if "config" in kwargs:
-                filename =  kwargs["config"]
-                if filename == None:
-                    filename = os.environ['KONFAI_CONFIG_FILE']
+                filename = kwargs["config"]
+                if filename is None:
+                    filename = os.environ["KONFAI_config_file"]
                 else:
-                    os.environ['KONFAI_CONFIG_FILE'] = filename
-                key_tmp =  kwargs["DL_args"]+("."+key if key is not None else "") if "DL_args" in kwargs else key
-                without =  kwargs["DL_without"] if "DL_without" in kwargs else []
-                os.environ['KONFAI_CONFIG_PATH'] = key_tmp
+                    os.environ["KONFAI_config_file"] = filename
+                key_tmp = (
+                    kwargs["konfai_args"] + ("." + key if key is not None else "") if "konfai_args" in kwargs else key
+                )
+                without = kwargs["konfai_without"] if "konfai_without" in kwargs else []
+                if key_tmp is not None:
+                    os.environ["KONFAI_CONFIG_PATH"] = key_tmp
                 with Config(filename, key_tmp) as config:
-                    os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
-                    kwargs = {} 
-                    for param in list(inspect.signature(function).parameters.values())[len(args):]:
+                    os.environ["KONFAI_CONFIG_VARIABLE"] = "False"
+                    kwargs = {}
+                    for param in list(inspect.signature(function).parameters.values())[len(args) :]:
                         if param.name in without:
                             continue
-                        
-                        annotation = param.annotation
 
+                        annotation = param.annotation
                         # --- support Literal ---
                         if get_origin(annotation) is Literal:
                             allowed_values = get_args(annotation)
                             default_value = param.default if param.default != inspect._empty else allowed_values[0]
-                            value = config.getValue(param.name, f"default:{default_value}")
+                            value = config.get_value(param.name, f"default:{default_value}")
                             if value not in allowed_values:
                                 raise ConfigError(
-                                    f"Invalid value '{value}' for parameter '{param.name} expected one of: {allowed_values}."
+                                    f"Invalid value '{value}' for parameter '{param.name} "
+                                    f"expected one of: {allowed_values}."
                                 )
                             kwargs[param.name] = value
                             continue
-                        if str(annotation).startswith("typing.Union") or str(annotation).startswith("typing.Optional"):
+
+                        if (
+                            str(annotation).startswith("typing.Union")
+                            or str(annotation).startswith("typing.Optional")
+                            or get_origin(annotation) is types.UnionType
+                        ):
                             for i in annotation.__args__:
                                 annotation = i
                                 break
-                        
+
                         if not annotation == inspect._empty:
                             if annotation not in [int, str, bool, float, torch.Tensor]:
-                                if str(annotation).startswith("list") or str(annotation).startswith("tuple") or str(annotation).startswith("typing.Tuple") or str(annotation).startswith("typing.List") or str(annotation).startswith("typing.Sequence"):
+                                if (
+                                    str(annotation).startswith("list")
+                                    or str(annotation).startswith("tuple")
+                                    or str(annotation).startswith("typing.Tuple")
+                                    or str(annotation).startswith("typing.List")
+                                    or str(annotation).startswith("typing.Sequence")
+                                ):
                                     elem_type = annotation.__args__[0]
-                                    values = config.getValue(param.name, param.default)
-                                    if getattr(elem_type, '__origin__', None) is Union:
+                                    values = config.get_value(param.name, param.default)
+                                    if getattr(elem_type, "__origin__", None) is Union:
                                         valid_types = elem_type.__args__
                                         result = []
                                         for v in values:
@@ -212,43 +242,70 @@ def config(key : Union[str, None] = None):
                                                     result.append(t(v) if t != torch.Tensor else v)
                                                     break
                                                 except Exception:
-                                                    continue
+                                                    raise ValueError("Merde")
                                         kwargs[param.name] = result
 
-                                    elif annotation.__args__[0] in [int, str, bool, float]:
-                                        values = config.getValue(param.name, param.default)
+                                    elif annotation.__args__[0] in [
+                                        int,
+                                        str,
+                                        bool,
+                                        float,
+                                    ]:
+                                        values = config.get_value(param.name, param.default)
                                         kwargs[param.name] = values
                                     else:
-                                        raise ConfigError("Config: The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]")
+                                        raise ConfigError(
+                                            "Config: The config only supports types : config(Object), int, str, bool,"
+                                            " float, list[int], list[str], list[bool], list[float], dict[str, Object]"
+                                        )
                                 elif str(annotation).startswith("dict"):
-                                    if annotation.__args__[0] == str:
-                                        values = config.getValue(param.name, param.default)
-                                        if values is not None and annotation.__args__[1] not in [int, str, bool, float, Any]:
+                                    if annotation.__args__[0] is str:
+                                        values = config.get_value(param.name, param.default)
+                                        if values is not None and annotation.__args__[1] not in [
+                                            int,
+                                            str,
+                                            bool,
+                                            float,
+                                            Any,
+                                        ]:
                                             try:
-                                                kwargs[param.name] = {value : annotation.__args__[1](config = filename, DL_args = key_tmp+"."+param.name+"."+value) for value in values}
+                                                kwargs[param.name] = {
+                                                    value: annotation.__args__[1](
+                                                        config=filename,
+                                                        konfai_args=str(key_tmp) + "." + param.name + "." + value,
+                                                    )
+                                                    for value in values
+                                                }
                                             except ValueError as e:
                                                 raise ValueError(e)
                                             except Exception as e:
-                                                raise ConfigError("{} {}".format(values, e))
+                                                raise ConfigError(f"{values} {e}")
                                         else:
 
                                             kwargs[param.name] = values
-                                    else: 
-                                        raise ConfigError("Config: The config only supports types : config(Object), int, str, bool, float, list[int], list[str], list[bool], list[float], dict[str, Object]")
+                                    else:
+                                        raise ConfigError(
+                                            "Config: The config only supports types : config(Object), int, str, bool,"
+                                            " float, list[int], list[str], list[bool], list[float], dict[str, Object]"
+                                        )
                                 else:
                                     try:
-                                        kwargs[param.name] = annotation(config = filename, DL_args = key_tmp)
+                                        kwargs[param.name] = annotation(config=filename, konfai_args=key_tmp)
                                     except Exception as e:
-                                        raise ConfigError("Failed to instantiate {} with type {}, error {} ".format(param.name, annotation, e))
-                                    
-                                    if os.environ['KONFAI_CONFIG_VARIABLE'] == "True":
-                                        os.environ['KONFAI_CONFIG_VARIABLE'] = "False"
+                                        raise ConfigError(
+                                            f"Failed to instantiate {param.name} with type {annotation}, error {e} "
+                                        )
+
+                                    if os.environ["KONFAI_CONFIG_VARIABLE"] == "True":
+                                        os.environ["KONFAI_CONFIG_VARIABLE"] = "False"
                                         kwargs[param.name] = None
                             else:
-                                kwargs[param.name] = config.getValue(param.name, param.default)
+                                kwargs[param.name] = config.get_value(param.name, param.default)
                         elif param.name != "self":
-                            kwargs[param.name] = config.getValue(param.name, param.default)
+                            kwargs[param.name] = config.get_value(param.name, param.default)
             result = function(*args, **kwargs)
             return result
+
         return new_function
+
     return decorator