From 4723f0d23d8be01792435a89c5370a6c27218a3a Mon Sep 17 00:00:00 2001
From: RGBCube <78925721+RGBCube@users.noreply.github.com>
Date: Mon, 27 Jun 2022 16:10:00 +0300
Subject: [PATCH] Cleanup unschema

---
 tools/unschema/__main__.py | 53 +++++++++++++++++++-------------------
 tools/unschema/parser.py   | 34 +++++++++++++++++++-----
 2 files changed, 55 insertions(+), 32 deletions(-)

diff --git a/tools/unschema/__main__.py b/tools/unschema/__main__.py
index 028628d..dd4fdf4 100644
--- a/tools/unschema/__main__.py
+++ b/tools/unschema/__main__.py
@@ -5,9 +5,21 @@ from argparse import ArgumentParser
 
 from .parser import generate
 
-parser = ArgumentParser(description="Generate TypedDicts from a json schema.")
-parser.add_argument("-f", "--from", required=True, help="The json schema to generate from.")
-parser.add_argument("-t", "--to", help="The file to write the TypedDicts to.")
+# fmt: off
+parser = ArgumentParser(
+    description="Generate TypedDicts from a json schema."
+)
+parser.add_argument(
+    "-f",
+    "--from",
+    required=True,
+    help="The json schema to generate from."
+)
+parser.add_argument(
+    "-t",
+    "--to",
+    help="The file to write the TypedDicts to."
+)
 parser.add_argument(
     "-on",
     "--object-name",
@@ -15,7 +27,10 @@ parser.add_argument(
     help="The name of the object to generate.",
 )
 parser.add_argument(
-    "-nc", "--no-comments", action="store_true", help="If given, comments wont be added."
+    "-nc",
+    "--no-comments",
+    action="store_true",
+    help="If given, comments wont be added."
 )
 parser.add_argument(
     "-p",
@@ -23,42 +38,28 @@ parser.add_argument(
     action="store_true",
     help="If given, the result will be printed.",
 )
+# fmt: on
 
 args = parser.parse_args()
 
 with open(args.__getattribute__("from")) as f:
-    generated = generate(json.load(f), no_comments=args.no_comments)
+    schema = json.load(f)
 
 start = time.perf_counter()
-
-text = f"""from __future__ import annotations
-
-from typing import Any, List, Optional, TypedDict, Union, TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing_extensions import NotRequired
-
-{generated[0]}
-
-{args.object_name} = {generated[1]}
-"""
+generated = generate(schema, no_comments=args.no_comments)
+end = time.perf_counter() - start
 
 if args.print:
-    print(text)
-
-    end = time.perf_counter() - start
-
+    print(generated)
 else:
     if not (to := args.to):
         raise ValueError("-t/--to is required when writing to a file.")
 
     with open(to, "w") as f:
-        f.write(text)
-
-    end = time.perf_counter() - start
+        f.write(generated)
 
     os.system(
-        f"unimport {to} --gitignore -r --ignore-init; isort {to}; black {to}; flynt {to} -tc",
+        f"unimport {to} --gitignore -r --ignore-init; black {to}",
     )
 
-print(f"\n\nSuccess! Finished in {end*1000:.3} milliseconds")
+print(f"\n\nSuccess! Finished in {end*1000:.3} milliseconds (formatting excluded)")
diff --git a/tools/unschema/parser.py b/tools/unschema/parser.py
index 5ac4890..cf62026 100644
--- a/tools/unschema/parser.py
+++ b/tools/unschema/parser.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+__all__ = ("generate",)
+
 from typing import Tuple  # , Any
 
 types = {
@@ -21,7 +23,7 @@ types = {
 #         super().append(obj)
 
 
-def generate(
+def _generate(
     obj: dict, /, *, title: str = "GeneratedObject", no_comments: bool = False
 ) -> Tuple[str, str]:  # sourcery skip: low-code-quality
     """Makes TypedDict from a JSON Schema object.
@@ -53,7 +55,7 @@ def generate(
                 if obj_schema["type"] == "null":
                     optional = True
                 else:
-                    extras, target = generate(obj_schema, title=title, no_comments=no_comments)
+                    extras, target = _generate(obj_schema, title=title, no_comments=no_comments)
                     result.append(extras)
                     union_items.append(target)
 
@@ -120,7 +122,7 @@ def generate(
             typed_dict = [f"class {obj_title}(TypedDict{total}):"]
 
             for key, value in obj_properties.items():
-                extras, param_annotation = generate(
+                extras, param_annotation = _generate(
                     value, title=key.capitalize(), no_comments=no_comments
                 )
                 result.append(extras)
@@ -156,7 +158,7 @@ def generate(
         # maxContains, minLength, maxLength, uniqueItems
 
         if obj_items := obj.get("items"):
-            extras, target = generate(obj_items, no_comments=no_comments)
+            extras, target = _generate(obj_items, no_comments=no_comments)
             result.append(extras)
             annotation = f"List[{target}]"
 
@@ -166,13 +168,13 @@ def generate(
                 (
                     extras,
                     target,
-                ) = generate(item, no_comments=no_comments)
+                ) = _generate(item, no_comments=no_comments)
                 result.append(extras)
                 tuple_annotation.append(target)
 
             if extra_item_type := obj.get("items"):
                 if extra_item_type is not True:
-                    extras, extra_type = generate(extra_item_type, no_comments=no_comments)
+                    extras, extra_type = _generate(extra_item_type, no_comments=no_comments)
                     result.append(extras)
                     if not no_comments:
                         result.append(
@@ -191,3 +193,23 @@ def generate(
     result = [i for i in result if i]
 
     return "\n".join(result), annotation
+
+
+text = """from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, List, Optional, TypedDict, Union
+
+if TYPE_CHECKING:
+    from typing_extensions import NotRequired
+
+{0[0]}
+
+{1} = {0[1]}
+"""
+
+
+def generate(
+    schema: dict, /, *, object_name: str = "GeneratedObject", no_comments: bool = False
+) -> str:
+    generated = _generate(schema, title=object_name, no_comments=no_comments)
+    return text.format(generated, object_name)