tensorneko-util 0.3.15__py3-none-any.whl → 0.3.17__py3-none-any.whl

tensorneko_util/io/json/json_writer.py

@@ -14,7 +14,7 @@ class JsonWriter:
                 json.dump(obj, file, indent=indent, ensure_ascii=ensure_ascii)
         else:
             import orjson
-            with open(path, "w", encoding=encoding) as file:
+            with open(path, "wb") as file:
                 if indent == 4:
                     warnings.warn("orjson does not support indent 4, will use indent 2 instead.")
                     option = orjson.OPT_INDENT_2

tensorneko_util/util/misc.py

@@ -241,7 +241,8 @@ def load_py(path: str) -> ModuleType:
     return module
 
 
-def try_until_success(func: Callable[[Any], T], *args, max_trials: Optional[int] = None, sleep_time: int = 0,
+def try_until_success(func: Callable[..., T], *args, max_trials: Optional[int] = None, sleep_time: int = 0,
+    exception_callback: Optional[Callable[[Exception], None]] = None,
     exception_type: Union[Type[T_E], Tuple[T_E, ...]] = Exception, **kwargs
 ) -> T:
     """
@@ -251,6 +252,7 @@ def try_until_success(func: Callable[[Any], T], *args, max_trials: Optional[int]
         func (``(...) -> T``): The function to run.
         max_trials (``int``, optional): The max try times. None for unlimited. Default: None.
         sleep_time (``int``, optional): The sleep time between each try. Default: 0.
+        exception_callback (``(Exception) -> None``, optional): The callback function when exception occurs.
         exception_type (``Type[Exception] | (Type[Exception], ...)``, optional): The exception types to catch.
             Default: Exception.
         *args: The args for the function.
@@ -268,6 +270,8 @@ def try_until_success(func: Callable[[Any], T], *args, max_trials: Optional[int]
             return func(*args, **kwargs)
         except exception_type as e:
             trials += 1
+            if exception_callback is not None:
+                exception_callback(e)
             if max_trials is not None and trials > max_trials:
                 raise e
         if sleep_time > 0:

tensorneko_util/util/registry.py

@@ -28,7 +28,7 @@ class Registry(Generic[T]):
     """
 
     def __init__(self):
-        self._registry: dict[str, T] = {}
+        self._registry: dict[str, Type[T]] = {}
 
     def register(self, name: str) -> Callable[[Type[T]], Type[T]]:
         def wrapper(cls):
@@ -37,7 +37,7 @@ class Registry(Generic[T]):
 
         return wrapper
 
-    def __getitem__(self, name: str) -> T:
+    def __getitem__(self, name: str) -> Type[T]:
         return self._registry[name]
 
     def items(self):

tensorneko_util/version.txt

@@ -1 +1 @@
-0.3.15
+0.3.17