add poly schemas support (#31)

* onward

* adding schemas for Pythonland!

* onward

* next

* next

* next

* next

* test

* next

* next

* next

* little tweak for A-Aron

* fix

* next

Author: eupharis

.flake8

@@ -1,4 +1,4 @@
 [flake8]
-extend-ignore = E203,E303,E402,E501,E722,W391,F401,W292
+ignore = E203,E303,E402,E501,E722,W391,F401,W292,F811
 max-line-length = 150
-max-complexity = 20
+max-complexity = 22

.gitignore

@@ -34,5 +34,6 @@ __pycache__
 .polyapi-python
 function_add_test.py
 lib_test*.py
-polyapi/poly/
-polyapi/vari/
+polyapi/poly
+polyapi/vari
+polyapi/schemas

README.md

@@ -143,6 +143,16 @@ To run this library's unit tests, please clone the repo then run:
 python -m unittest discover
 ```
 
+## Linting
+
+The flake8 config is at the root of this repo at `.flake8`.
+
+When hacking on this library, please enable flake8 and add this line to your flake8 args (e.g., in your VSCode Workspace Settings):
+
+```
+--config=.flake8
+```
+
 ## Support
 
 If you run into any issues or want help getting started with this project, please contact support@polyapi.io

polyapi/api.py

@@ -8,6 +8,7 @@
 from typing import List, Dict, Any, TypedDict
 {args_def}
 {return_type_def}
+
 class {api_response_type}(TypedDict):
     status: int
     headers: Dict

polyapi/cli.py

@@ -13,6 +13,7 @@
 
 CLI_COMMANDS = ["setup", "generate", "function", "clear", "help", "update_rendered_spec"]
 
+
 def execute_from_cli():
     # First we setup all our argument parsing logic
     # Then we parse the arguments (waaay at the bottom)
@@ -46,7 +47,7 @@ def setup(args):
 
     def generate_command(args):
         initialize_config()
-        print("Generating Poly functions...", end="")
+        print("Generating Poly Python SDK...", end="")
         generate()
         print_green("DONE")
 

polyapi/function_cli.py

@@ -1,7 +1,7 @@
 import sys
 from typing import Any, List, Optional
 import requests
-from polyapi.generate import get_functions_and_parse, generate_functions
+from polyapi.generate import cache_specs, generate_functions, get_specs, parse_function_specs
 from polyapi.config import get_api_key_and_url
 from polyapi.utils import get_auth_headers, print_green, print_red, print_yellow
 from polyapi.parser import parse_function_code, get_jsonschema_type
@@ -88,7 +88,10 @@ def function_add_or_update(
         print(f"Function ID: {function_id}")
         if generate:
             print("Generating new custom function...", end="")
-            functions = get_functions_and_parse(limit_ids=[function_id])
+            # TODO do something more efficient here rather than regetting ALL the specs again
+            specs = get_specs()
+            cache_specs(specs)
+            functions = parse_function_specs(specs)
             generate_functions(functions)
             print_green("DONE")
     else:

polyapi/generate.py

@@ -2,13 +2,15 @@
 import requests
 import os
 import shutil
-from typing import List
+from typing import List, cast
 
+from polyapi import schema
 from polyapi.auth import render_auth_function
 from polyapi.client import render_client_function
+from polyapi.poly_schemas import generate_schemas
 from polyapi.webhook import render_webhook_handle
 
-from .typedefs import PropertySpecification, SpecificationDto, VariableSpecDto
+from .typedefs import PropertySpecification, SchemaSpecDto, SpecificationDto, VariableSpecDto
 from .api import render_api_function
 from .server import render_server_function
 from .utils import add_import_to_init, get_auth_headers, init_the_init, to_func_namespace
@@ -18,12 +20,21 @@
 SUPPORTED_FUNCTION_TYPES = {
     "apiFunction",
     "authFunction",
-    "customFunction",
+    "customFunction",  # client function - this is badly named in /specs atm
     "serverFunction",
     "webhookHandle",
 }
 
-SUPPORTED_TYPES = SUPPORTED_FUNCTION_TYPES | {"serverVariable"}
+SUPPORTED_TYPES = SUPPORTED_FUNCTION_TYPES | {"serverVariable", "schema", "snippet"}
+
+
+X_POLY_REF_WARNING = '''"""
+x-poly-ref:
+  path:'''
+
+X_POLY_REF_BETTER_WARNING = '''"""
+Unresolved schema, please add the following schema to complete it:
+  path:'''
 
 
 def get_specs() -> List:
@@ -38,9 +49,56 @@ def get_specs() -> List:
         raise NotImplementedError(resp.content)
 
 
+def build_schema_index(items):
+    index = {}
+    for item in items:
+        if item.get("type") == "schema" and "contextName" in item:
+            index[item["contextName"]] = {**item.get("definition", {}), "name": item.get("name")}
+    return index
+
+
+def resolve_poly_refs(obj, schema_index):
+    if isinstance(obj, dict):
+        if "x-poly-ref" in obj:
+            ref = obj["x-poly-ref"]
+            if isinstance(ref, dict) and "path" in ref:
+                path = ref["path"]
+                if path in schema_index:
+                    return resolve_poly_refs(schema_index[path], schema_index)
+                else:
+                    return obj
+        return {k: resolve_poly_refs(v, schema_index) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [resolve_poly_refs(item, schema_index) for item in obj]
+    else:
+        return obj
+
+
+def replace_poly_refs_in_functions(specs: List[SpecificationDto], schema_index):
+    spec_idxs_to_remove = []
+    for idx, spec in enumerate(specs):
+        if spec.get("type") in ("apiFunction", "customFunction", "serverFunction"):
+            func = spec.get("function")
+            if func:
+                try:
+                    spec["function"] = resolve_poly_refs(func, schema_index)
+                except Exception:
+                    print()
+                    print(f"{spec['context']}.{spec['name']} (id: {spec['id']}) failed to resolve poly refs, skipping!")
+                    spec_idxs_to_remove.append(idx)
+
+    # reverse the list so we pop off later indexes first
+    spec_idxs_to_remove.reverse()
+
+    for idx in spec_idxs_to_remove:
+        specs.pop(idx)
+
+    return specs
+
+
 def parse_function_specs(
     specs: List[SpecificationDto],
-    limit_ids: List[str] | None,  # optional list of ids to limit to
+    limit_ids: List[str] | None = None,  # optional list of ids to limit to
 ) -> List[SpecificationDto]:
     functions = []
     for spec in specs:
@@ -91,23 +149,14 @@ def read_cached_specs() -> List[SpecificationDto]:
         return json.loads(f.read())
 
 
-def get_functions_and_parse(limit_ids: List[str] | None = None) -> List[SpecificationDto]:
-    specs = get_specs()
-    cache_specs(specs)
-    return parse_function_specs(specs, limit_ids=limit_ids)
+def get_variables() -> List[VariableSpecDto]:
+    specs = read_cached_specs()
+    return [cast(VariableSpecDto, spec) for spec in specs if spec["type"] == "serverVariable"]
 
 
-def get_variables() -> List[VariableSpecDto]:
-    api_key, api_url = get_api_key_and_url()
-    headers = {"Authorization": f"Bearer {api_key}"}
-    # TODO do some caching so this and get_functions just do 1 function call
-    url = f"{api_url}/specs"
-    resp = requests.get(url, headers=headers)
-    if resp.status_code == 200:
-        specs = resp.json()
-        return [spec for spec in specs if spec["type"] == "serverVariable"]
-    else:
-        raise NotImplementedError(resp.content)
+def get_schemas() -> List[SchemaSpecDto]:
+    specs = read_cached_specs()
+    return [cast(SchemaSpecDto, spec) for spec in specs if spec["type"] == "schema"]
 
 
 def remove_old_library():
@@ -120,12 +169,28 @@ def remove_old_library():
     if os.path.exists(path):
         shutil.rmtree(path)
 
+    path = os.path.join(currdir, "schemas")
+    if os.path.exists(path):
+        shutil.rmtree(path)
+
 
 def generate() -> None:
 
     remove_old_library()
 
-    functions = get_functions_and_parse()
+    limit_ids: List[str] = []  # useful for narrowing down generation to a single function to debug
+
+    specs = get_specs()
+    cache_specs(specs)
+    functions = parse_function_specs(specs, limit_ids=limit_ids)
+
+    schemas = get_schemas()
+    if schemas:
+        generate_schemas(schemas)
+
+    schema_index = build_schema_index(schemas)
+    functions = replace_poly_refs_in_functions(functions, schema_index)
+
     if functions:
         generate_functions(functions)
     else:
@@ -138,6 +203,7 @@ def generate() -> None:
     if variables:
         generate_variables(variables)
 
+
     # indicator to vscode extension that this is a polyapi-python project
     file_path = os.path.join(os.getcwd(), ".polyapi-python")
     open(file_path, "w").close()
@@ -214,6 +280,12 @@ def render_spec(spec: SpecificationDto):
             arguments,
             return_type,
         )
+
+    if X_POLY_REF_WARNING in func_type_defs:
+        # this indicates that jsonschema_gentypes has detected an x-poly-ref
+        # let's add a more user friendly error explaining what is going on
+        func_type_defs = func_type_defs.replace(X_POLY_REF_WARNING, X_POLY_REF_BETTER_WARNING)
+
     return func_str, func_type_defs
 
 

polyapi/poly/__init__.py

@@ -0,0 +1,61 @@
+import os
+from typing import Any, Dict, List, Tuple
+
+from polyapi.schema import wrapped_generate_schema_types
+from polyapi.utils import add_import_to_init, init_the_init
+from tests.test_schema import SCHEMA
+
+from .typedefs import SchemaSpecDto
+
+SCHEMA_CODE_IMPORTS = """from typing_extensions import TypedDict, NotRequired
+
+
+"""
+
+
+FALLBACK_SPEC_TEMPLATE = """class {name}(TypedDict, total=False):
+    ''' unable to generate schema for {name}, defaulting to permissive type '''
+    pass
+"""
+
+
+def generate_schemas(specs: List[SchemaSpecDto]):
+    for spec in specs:
+        create_schema(spec)
+
+
+def create_schema(spec: SchemaSpecDto) -> None:
+    folders = ["schemas"]
+    if spec["context"]:
+        folders += [s for s in spec["context"].split(".")]
+
+    # build up the full_path by adding all the folders
+    full_path = os.path.join(os.path.dirname(os.path.abspath(__file__)))
+
+    for idx, folder in enumerate(folders):
+        full_path = os.path.join(full_path, folder)
+        if not os.path.exists(full_path):
+            os.makedirs(full_path)
+        next = folders[idx + 1] if idx + 1 < len(folders) else None
+        if next:
+            add_import_to_init(full_path, next, code_imports=SCHEMA_CODE_IMPORTS)
+
+    add_schema_to_init(full_path, spec)
+
+
+def add_schema_to_init(full_path: str, spec: SchemaSpecDto):
+    init_the_init(full_path, code_imports="")
+    init_path = os.path.join(full_path, "__init__.py")
+    with open(init_path, "a") as f:
+        f.write(render_poly_schema(spec) + "\n\n")
+
+
+def render_poly_schema(spec: SchemaSpecDto) -> str:
+    definition = spec["definition"]
+    if not definition.get("type"):
+        definition["type"] = "object"
+    root, schema_types = wrapped_generate_schema_types(
+        definition, root=spec["name"], fallback_type=Dict
+    )
+    return schema_types
+    # return FALLBACK_SPEC_TEMPLATE.format(name=spec["name"])

polyapi/poly_schemas.py

@@ -1,11 +1,18 @@
+""" NOTE: this file represents the schema parsing logic for jsonschema_gentypes
+"""
+import random
+import string
 import logging
 import contextlib
 from typing import Dict
 from jsonschema_gentypes.cli import process_config
 from jsonschema_gentypes import configuration
+import referencing
 import tempfile
 import json
 
+import referencing.exceptions
+
 from polyapi.constants import JSONSCHEMA_TO_PYTHON_TYPE_MAP
 
 
@@ -33,8 +40,18 @@ def _temp_store_input_data(input_data: Dict) -> str:
 
 
 def wrapped_generate_schema_types(type_spec: dict, root, fallback_type):
+    from polyapi.utils import pascalCase
     if not root:
-        root = "MyList" if fallback_type == "List" else "MyDict"
+        root = "List" if fallback_type == "List" else "Dict"
+        if type_spec.get("x-poly-ref") and type_spec["x-poly-ref"].get("path"):
+            # x-poly-ref occurs when we have an unresolved reference
+            # lets name the root after the reference for some level of visibility
+            root += pascalCase(type_spec["x-poly-ref"]["path"].replace(".", " "))
+        else:
+            # add three random letters for uniqueness
+            root += random.choice(string.ascii_letters).upper()
+            root += random.choice(string.ascii_letters).upper()
+            root += random.choice(string.ascii_letters).upper()
 
     root = clean_title(root)
 
@@ -44,8 +61,13 @@ def wrapped_generate_schema_types(type_spec: dict, root, fallback_type):
         # some schemas are so huge, our library cant handle it
         # TODO identify critical recursion penalty and maybe switch underlying logic to iterative?
         return fallback_type, ""
+    except referencing.exceptions.CannotDetermineSpecification:
+        # just go with fallback_type here
+        # we couldn't match the right $ref earlier in resolve_poly_refs
+        # {'$ref': '#/definitions/FinanceAccountListModel'}
+        return fallback_type, ""
     except:
-        logging.exception(f"Error when generating schema type: {type_spec}")
+        logging.error(f"Error when generating schema type: {type_spec}\nusing fallback type '{fallback_type}'")
         return fallback_type, ""