152 lines
5.0 KiB
Python
152 lines
5.0 KiB
Python
import argparse
|
|
import json
|
|
from typing import Any
|
|
|
|
|
|
# Function to map JSON schema types to Python types
|
|
def map_json_type(json_type: Any, nested_type: str = "Any") -> str:
|
|
if isinstance(json_type, list):
|
|
return " | ".join(map(map_json_type, json_type))
|
|
if isinstance(json_type, dict):
|
|
return map_json_type(json_type.get("type"))
|
|
elif json_type == "string":
|
|
return "str"
|
|
elif json_type == "integer":
|
|
return "int"
|
|
elif json_type == "boolean":
|
|
return "bool"
|
|
elif json_type == "array":
|
|
return f"list[{nested_type}]" # Further specification can be handled if needed
|
|
elif json_type == "object":
|
|
return f"dict[str, {nested_type}]"
|
|
elif json_type == "null":
|
|
return "None"
|
|
else:
|
|
return "Any"
|
|
|
|
|
|
known_classes = set()
|
|
root_class = "Inventory"
|
|
|
|
|
|
# Recursive function to generate dataclasses from JSON schema
|
|
def generate_dataclass(schema: dict[str, Any], class_name: str = root_class) -> str:
|
|
properties = schema.get("properties", {})
|
|
required = schema.get("required", [])
|
|
|
|
fields = []
|
|
nested_classes = []
|
|
|
|
for prop, prop_info in properties.items():
|
|
field_name = prop.replace("-", "_")
|
|
|
|
prop_type = prop_info.get("type", None)
|
|
union_variants = prop_info.get("oneOf", [])
|
|
|
|
title = prop_info.get("title", prop.removesuffix("s"))
|
|
title_sanitized = "".join([p.capitalize() for p in title.split("-")])
|
|
nested_class_name = f"""{class_name if class_name != root_class and not prop_info.get("title") else ""}{title_sanitized}"""
|
|
|
|
# if nested_class_name == "ServiceBorgbackupRoleServerConfig":
|
|
# breakpoint()
|
|
|
|
if (prop_type is None) and (not union_variants):
|
|
raise ValueError(f"Type not found for property {prop} {prop_info}")
|
|
|
|
# Unions fields (oneOf)
|
|
# str | int | None
|
|
python_type = None
|
|
|
|
if union_variants:
|
|
python_type = map_json_type(union_variants)
|
|
elif prop_type == "array":
|
|
item_schema = prop_info.get("items")
|
|
if isinstance(item_schema, dict):
|
|
python_type = map_json_type(
|
|
prop_type,
|
|
map_json_type(item_schema),
|
|
)
|
|
else:
|
|
python_type = map_json_type(
|
|
prop_type,
|
|
map_json_type([i for i in prop_info.get("items", [])]),
|
|
)
|
|
|
|
assert python_type, f"Python type not found for {prop} {prop_info}"
|
|
|
|
if prop in required:
|
|
field_def = f"{prop}: {python_type}"
|
|
else:
|
|
field_def = f"{prop}: {python_type} | None = None"
|
|
|
|
if prop_type == "object":
|
|
map_type = prop_info.get("additionalProperties")
|
|
if map_type:
|
|
# breakpoint()
|
|
if map_type.get("type") == "object":
|
|
# Non trivial type
|
|
if nested_class_name not in known_classes:
|
|
nested_classes.append(
|
|
generate_dataclass(map_type, nested_class_name)
|
|
)
|
|
known_classes.add(nested_class_name)
|
|
field_def = f"{field_name}: dict[str, {nested_class_name}]"
|
|
else:
|
|
# Trivial type
|
|
field_def = f"{field_name}: dict[str, {map_json_type(map_type)}]"
|
|
else:
|
|
if nested_class_name not in known_classes:
|
|
nested_classes.append(
|
|
generate_dataclass(prop_info, nested_class_name)
|
|
)
|
|
known_classes.add(nested_class_name)
|
|
|
|
field_def = f"{field_name}: {nested_class_name}"
|
|
|
|
elif prop_type == "array":
|
|
items = prop_info.get("items", {})
|
|
if items.get("type") == "object":
|
|
nested_class_name = prop.capitalize()
|
|
nested_classes.append(generate_dataclass(items, nested_class_name))
|
|
field_def = f"{field_name}: List[{nested_class_name}]"
|
|
|
|
fields.append(field_def)
|
|
|
|
fields_str = "\n ".join(fields)
|
|
nested_classes_str = "\n\n".join(nested_classes)
|
|
|
|
class_def = f"@dataclass\nclass {class_name}:\n {fields_str}\n"
|
|
return f"{nested_classes_str}\n\n{class_def}" if nested_classes_str else class_def
|
|
|
|
|
|
def run_gen(args: argparse.Namespace) -> None:
|
|
print(f"Converting {args.input} to {args.output}")
|
|
dataclass_code = ""
|
|
with open(args.input) as f:
|
|
schema = json.load(f)
|
|
dataclass_code = generate_dataclass(schema)
|
|
|
|
with open(args.output, "w") as f:
|
|
f.write(
|
|
"""
|
|
from dataclasses import dataclass
|
|
from typing import Any\n\n
|
|
"""
|
|
)
|
|
f.write(dataclass_code)
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("input", help="Input JSON schema file")
|
|
parser.add_argument("output", help="Output Python file")
|
|
parser.set_defaults(func=run_gen)
|
|
|
|
args = parser.parse_args()
|
|
|
|
args.func(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|