tensorneko-util 0.3.12__tar.gz → 0.3.14__tar.gz

{tensorneko_util-0.3.12 → tensorneko_util-0.3.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tensorneko_util
-Version: 0.3.12
+Version: 0.3.14
 Summary: The Utils for Library TensorNeko.
 Home-page: https://github.com/ControlNet/tensorneko
 Author: ControlNet

{tensorneko_util-0.3.12 → tensorneko_util-0.3.14}/src/tensorneko_util/io/json/json_reader.py

@@ -1,4 +1,5 @@
 import json
+import warnings
 from typing import Union, Optional, List
 
 from ...util.type import T
@@ -6,8 +7,19 @@ from ...util.type import T
 
 class JsonReader:
 
+    @staticmethod
+    def _read_json(path: str, encoding: str = "UTF-8", fast: bool = False) -> Union[dict, list]:
+        if not fast:
+            with open(path, "r", encoding=encoding) as file:
+                return json.load(file)
+        else:
+            import orjson
+            with open(path, "rb") as file:
+                return orjson.loads(file.read())
+
     @classmethod
-    def of(cls, path: str, clazz: Optional[T] = None, encoding: str = "UTF-8") -> Union[T, dict, list]:
+    def of(cls, path: str, clazz: Optional[T] = None, encoding: str = "UTF-8", fast: bool = True
+    ) -> Union[T, dict, list]:
         """
         Read json files to ``list`` or ``dict``.
 
@@ -16,14 +28,23 @@ class JsonReader:
             clazz: (``T``, optional): The object of the json read for. The type should be decorated by
                 :func:`json_data`. This should be ``T`` or ``List[T]`` or ``List[List[T]]``
             encoding (``str``, optional): The encoding for python ``open`` function. Default: "UTF-8"
+            fast (``bool``, optional): Whether to use faster `orjson`. If `orjson` is not installed, use
+                `json` library. Default: True
 
         Returns:
             ``dict`` | ``list`` | ``object``: The object of given json.
         """
+        if fast:
+            try:
+                import orjson
+            except ImportError:
+                fast = False
+                warnings.warn("orjson is not installed, will use json lib instead.")
+            assert encoding == "UTF-8", "orjson only supports UTF-8 encoding."
+
         if clazz is not None:
             if "typing.List[" in str(clazz):
-                with open(path, "r", encoding=encoding) as file:
-                    obj: list = json.loads(file.read())
+                obj = cls._read_json(path, encoding, fast)
 
                 inner_type = clazz.__args__[0]
                 if "typing.List[" in str(inner_type):
@@ -44,17 +65,16 @@ class JsonReader:
                     if is_json_data:
                         obj = list(map(inner_type, obj))
             else:
-                with open(path, "r", encoding=encoding) as file:
-                    obj = clazz(json.loads(file.read()))
+                obj = clazz(cls._read_json(path, encoding, fast))
         else:
-            with open(path, "r", encoding=encoding) as file:
-                obj = json.loads(file.read())
+            obj = cls._read_json(path, encoding, fast)
 
         return obj
 
 
     @classmethod
-    def of_jsonl(cls, path: str, clazz: Optional[T] = None, encoding: str = "UTF-8") -> List[Union[T, dict, list]]:
+    def of_jsonl(cls, path: str, clazz: Optional[T] = None, encoding: str = "UTF-8", fast: bool = True
+    ) -> List[Union[T, dict, list]]:
         """
         Read jsonl files to ``list`` or ``dict``.
 
@@ -63,14 +83,27 @@ class JsonReader:
             clazz: (``T``, optional): The object of the jsonl read for. The type should be decorated by
                 :func:`json_data`. This should be ``T`` or ``List[T]`` or ``List[List[T]]``
             encoding (``str``, optional): The encoding for python ``open`` function. Default: "UTF-8"
+            fast (``bool``, optional): Whether to use faster `orjson`. If `orjson` is not installed, use
+                `json` library. Default: True
 
         Returns:
             ``List[dict]`` | ``List[list]`` | ``List[T]``: The object of given jsonl.
         """
+        if fast:
+            try:
+                import orjson
+            except ImportError:
+                fast = False
+                warnings.warn("orjson is not installed, will use json lib instead.")
+            assert encoding == "UTF-8", "orjson only supports UTF-8 encoding."
+
         if clazz is not None:
             if "typing.List[" in str(clazz):
                 with open(path, "r", encoding=encoding) as file:
-                    obj: list = [json.loads(line) for line in file]
+                    if fast:
+                        obj: list = [orjson.loads(line) for line in file]
+                    else:
+                        obj: list = [json.loads(line) for line in file]
 
                 inner_type = clazz.__args__[0]
                 if "typing.List[" in str(inner_type):
@@ -92,14 +125,21 @@ class JsonReader:
                         obj = list(map(inner_type, obj))
             else:
                 with open(path, "r", encoding=encoding) as file:
-                    obj = clazz([json.loads(line) for line in file])
+                    if fast:
+                        obj: list = [orjson.loads(line) for line in file]
+                    else:
+                        obj: list = [json.loads(line) for line in file]
+                    obj = clazz(obj)
         else:
             with open(path, "r", encoding=encoding) as file:
-                obj = [json.loads(line) for line in file]
+                if fast:
+                    obj: list = [orjson.loads(line) for line in file]
+                else:
+                    obj: list = [json.loads(line) for line in file]
 
         return obj
 
-    def __new__(cls, path: str, clazz: T = None, encoding: str = "UTF-8") -> Union[T, dict, list]:
+    def __new__(cls, path: str, clazz: T = None, encoding: str = "UTF-8", fast: bool = True) -> Union[T, dict, list]:
         """
         Read json or jsonl file smartly.
 
@@ -108,11 +148,13 @@ class JsonReader:
             clazz: (``T``, optional): The object of the json read for. The type should be decorated by
                 :func:`json_data`. This should be ``T`` or ``List[T]`` or ``List[List[T]]``
             encoding (``str``, optional): The encoding for python ``open`` function. Default: "UTF-8"
+            fast (``bool``, optional): Whether to use faster `orjson`. If `orjson` is not installed, use
+                `json` library. Default: True
 
         Returns:
             ``dict`` | ``list`` | ``object``: The object of given json or jsonl.
         """
         if path.endswith(".jsonl"):
-            return cls.of_jsonl(path, clazz, encoding)
+            return cls.of_jsonl(path, clazz, encoding, fast)
         else:
-            return cls.of(path, clazz, encoding)
+            return cls.of(path, clazz, encoding, fast)

tensorneko_util-0.3.14/src/tensorneko_util/io/json/json_writer.py

@@ -0,0 +1,70 @@
+import json
+import warnings
+from typing import Union
+
+
+class JsonWriter:
+
+    @staticmethod
+    def _write_json(path: str, obj: Union[dict, list], encoding: str = "UTF-8", indent: int = 4,
+        ensure_ascii: bool = False, fast: bool = True
+    ) -> None:
+        if not fast:
+            with open(path, "w", encoding=encoding) as file:
+                json.dump(obj, file, indent=indent, ensure_ascii=ensure_ascii)
+        else:
+            import orjson
+            with open(path, "w", encoding=encoding) as file:
+                if indent == 4:
+                    option = orjson.OPT_INDENT_4
+                elif indent == 2:
+                    option = orjson.OPT_INDENT_2
+                else:
+                    option = None
+                file.write(orjson.dumps(obj, option=option))
+
+    @classmethod
+    def to(cls, path: str, obj: Union[dict, list, object], encoding: str = "UTF-8", indent: int = 4,
+        ensure_ascii: bool = False, fast: bool = True
+    ) -> None:
+        """
+        Save as Json file from a dictionary, list or json_dict.
+
+        Args:
+            path (``str``): The path of output file.
+            obj (``dict`` | ``list`` | ``object``): The json data which need to be used for output. The type should be
+                ``dict``, ``list`` or the ``object`` decorated by :func:`~tensorneko.io.text.text_reader.json_data`.
+            encoding (``str``, optional): Python file IO encoding parameter. Default: "UTF-8".
+            indent (``int``, optional): Python json dump indent parameter. Default: 4.
+            ensure_ascii (``bool``, optional): Python json dump ensure_ascii parameter. Default: False.
+            fast (``bool``, optional): Whether to use faster `orjson`. If `orjson` is not installed, use
+                `json` library. Default: True
+        """
+        if fast:
+            try:
+                import orjson
+            except ImportError:
+                fast = False
+                warnings.warn("orjson is not installed, will use json lib instead.")
+            assert encoding == "UTF-8", "orjson only supports UTF-8 encoding."
+            assert indent in (0, 2, 4, None), "Only support indent 2, 4 or 0/None."
+            assert ensure_ascii is False, "orjson does not support ensure_ascii."
+
+        if type(obj) is dict:
+            cls._write_json(path, obj, encoding, indent, ensure_ascii, fast)
+        elif "is_json_data" in type(obj).__dict__ and type(obj).is_json_data:
+            cls._write_json(path, obj.to_dict(), encoding, indent, ensure_ascii, fast)
+        elif type(obj) is list:
+            if len(obj) == 0 or type(obj[0]) in (dict, list):
+                cls._write_json(path, obj, encoding, indent, ensure_ascii, fast)
+            elif "is_json_data" in type(obj[0]).__dict__ and type(obj[0]).is_json_data:
+                obj = [each.to_dict() for each in obj]
+                cls._write_json(path, obj, encoding, indent, ensure_ascii, fast)
+        else:
+            raise TypeError("Not implemented type. Only support dict, list, json_data.")
+
+    def __new__(cls, path: str, obj: Union[dict, list, object], encoding: str = "UTF-8", indent: int = 4,
+        ensure_ascii: bool = False, fast: bool = True
+    ) -> None:
+        """Alias of :meth:`~tensorneko.io.json.json_writer.JsonWriter.to`."""
+        cls.to(path, obj, encoding, indent, ensure_ascii, fast)

tensorneko_util-0.3.14/src/tensorneko_util/version.txt

@@ -0,0 +1 @@
+0.3.14