77 lines
2.7 KiB
Python
77 lines
2.7 KiB
Python
from __future__ import annotations
|
|
from copy import deepcopy
|
|
from typing import List, Optional
|
|
|
|
|
|
class SchemaProcessor:
|
|
|
|
def resolve_refs(self, schema, defs):
|
|
if isinstance(schema, dict):
|
|
if "$ref" in schema:
|
|
ref_path = schema["$ref"]
|
|
if ref_path.startswith("#/$defs/"):
|
|
def_key = ref_path.replace("#/$defs/", "")
|
|
return self.resolve_refs(defs[def_key], defs)
|
|
else:
|
|
raise ValueError(f"Unsupported $ref path: {ref_path}")
|
|
else:
|
|
return {k: self.resolve_refs(v, defs) for k, v in schema.items()}
|
|
elif isinstance(schema, list):
|
|
return [self.resolve_refs(item, defs) for item in schema]
|
|
else:
|
|
return schema
|
|
|
|
def flatten_schema(self, schema):
|
|
schema = deepcopy(schema)
|
|
defs = schema.pop("$defs", {})
|
|
return self.resolve_refs(schema, defs)
|
|
|
|
def find_dict_with_key(
|
|
self, schema: dict, target_key: str, current_path: Optional[List[str]] = None
|
|
) -> List[List[str]]:
|
|
if current_path is None:
|
|
current_path = []
|
|
paths = []
|
|
if target_key in schema:
|
|
paths.append(current_path.copy())
|
|
|
|
for key, value in schema.items():
|
|
if isinstance(value, dict):
|
|
new_path = current_path + [key]
|
|
paths.extend(self.find_dict_with_key(value, target_key, new_path))
|
|
elif isinstance(value, list):
|
|
for i, item in enumerate(value):
|
|
if isinstance(item, dict):
|
|
new_path = current_path + [key, str(i)]
|
|
paths.extend(
|
|
self.find_dict_with_key(item, target_key, new_path)
|
|
)
|
|
return paths
|
|
|
|
def get_dict_at_path(self, schema: dict, path: List[str]) -> dict:
|
|
current = schema
|
|
|
|
for part in path:
|
|
if part.isdigit():
|
|
current = current[int(part)]
|
|
else:
|
|
current = current[part]
|
|
|
|
return current
|
|
|
|
def remove_image_url_fields(self, schema: dict) -> dict:
|
|
copied_schema = schema.copy()
|
|
|
|
image_type_paths = self.find_dict_with_key(copied_schema, "_image_type")
|
|
|
|
for path in image_type_paths:
|
|
dict_at_path = self.get_dict_at_path(copied_schema, path)
|
|
|
|
if dict_at_path.get("_image_type") == "image":
|
|
if "properties" in dict_at_path and "url" in dict_at_path["properties"]:
|
|
del dict_at_path["properties"]["url"]
|
|
|
|
if "required" in dict_at_path and "url" in dict_at_path["required"]:
|
|
dict_at_path["required"].remove("url")
|
|
|
|
return copied_schema
|