Skip to content

Commit 54da596

Browse files
bokelleyclaude
andauthored
feat: improve type generation with discriminated union support (#21)
- Pull latest AdCP schemas (v1.0.0) with new destination and deployment types - Add discriminated union generation for oneOf schemas with type discriminators - Handle const fields for proper Literal types in discriminators - Add missing ActivationKey type alias for deployment schema Changes to type generation: - New `generate_discriminated_union()` function handles oneOf schemas - Automatically detects discriminator field (type) and generates typed variants - Example: Destination = PlatformDestination | AgentDestination - Each variant gets proper Literal["platform"] | Literal["agent"] type New types added: - PlatformDestination / AgentDestination (discriminated union) - PlatformDeployment / AgentDeployment (discriminated union) - ActivationKey type alias Improved existing types: - BrandManifestRef now discriminated union (was Any) - StartTiming now discriminated union (was Any) - PricingOption now discriminated union (was Any) All tests pass with Python 3.11+. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude <noreply@anthropic.com>
1 parent 1944b39 commit 54da596

File tree

4 files changed

+368
-6
lines changed

4 files changed

+368
-6
lines changed
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
{
2+
"$schema": "http://json-schema.org/draft-07/schema#",
3+
"$id": "/schemas/v1/core/deployment.json",
4+
"title": "Deployment",
5+
"description": "A signal deployment to a specific destination platform with activation status and key",
6+
"oneOf": [
7+
{
8+
"type": "object",
9+
"properties": {
10+
"type": {
11+
"const": "platform",
12+
"description": "Discriminator indicating this is a platform-based deployment"
13+
},
14+
"platform": {
15+
"type": "string",
16+
"description": "Platform identifier for DSPs"
17+
},
18+
"account": {
19+
"type": "string",
20+
"description": "Account identifier if applicable"
21+
},
22+
"is_live": {
23+
"type": "boolean",
24+
"description": "Whether signal is currently active on this destination"
25+
},
26+
"activation_key": {
27+
"$ref": "activation-key.json",
28+
"description": "The key to use for targeting. Only present if is_live=true AND requester has access to this destination."
29+
},
30+
"estimated_activation_duration_minutes": {
31+
"type": "number",
32+
"description": "Estimated time to activate if not live, or to complete activation if in progress",
33+
"minimum": 0
34+
},
35+
"deployed_at": {
36+
"type": "string",
37+
"format": "date-time",
38+
"description": "Timestamp when activation completed (if is_live=true)"
39+
}
40+
},
41+
"required": [
42+
"type",
43+
"platform",
44+
"is_live"
45+
],
46+
"additionalProperties": false
47+
},
48+
{
49+
"type": "object",
50+
"properties": {
51+
"type": {
52+
"const": "agent",
53+
"description": "Discriminator indicating this is an agent URL-based deployment"
54+
},
55+
"agent_url": {
56+
"type": "string",
57+
"format": "uri",
58+
"description": "URL identifying the destination agent"
59+
},
60+
"account": {
61+
"type": "string",
62+
"description": "Account identifier if applicable"
63+
},
64+
"is_live": {
65+
"type": "boolean",
66+
"description": "Whether signal is currently active on this destination"
67+
},
68+
"activation_key": {
69+
"$ref": "activation-key.json",
70+
"description": "The key to use for targeting. Only present if is_live=true AND requester has access to this destination."
71+
},
72+
"estimated_activation_duration_minutes": {
73+
"type": "number",
74+
"description": "Estimated time to activate if not live, or to complete activation if in progress",
75+
"minimum": 0
76+
},
77+
"deployed_at": {
78+
"type": "string",
79+
"format": "date-time",
80+
"description": "Timestamp when activation completed (if is_live=true)"
81+
}
82+
},
83+
"required": [
84+
"type",
85+
"agent_url",
86+
"is_live"
87+
],
88+
"additionalProperties": false
89+
}
90+
]
91+
}
Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
{
2+
"$schema": "http://json-schema.org/draft-07/schema#",
3+
"$id": "/schemas/v1/core/destination.json",
4+
"title": "Destination",
5+
"description": "A destination platform where signals can be activated (DSP, sales agent, etc.)",
6+
"oneOf": [
7+
{
8+
"type": "object",
9+
"properties": {
10+
"type": {
11+
"const": "platform",
12+
"description": "Discriminator indicating this is a platform-based destination"
13+
},
14+
"platform": {
15+
"type": "string",
16+
"description": "Platform identifier for DSPs (e.g., 'the-trade-desk', 'amazon-dsp')"
17+
},
18+
"account": {
19+
"type": "string",
20+
"description": "Optional account identifier on the platform"
21+
}
22+
},
23+
"required": [
24+
"type",
25+
"platform"
26+
],
27+
"additionalProperties": false
28+
},
29+
{
30+
"type": "object",
31+
"properties": {
32+
"type": {
33+
"const": "agent",
34+
"description": "Discriminator indicating this is an agent URL-based destination"
35+
},
36+
"agent_url": {
37+
"type": "string",
38+
"format": "uri",
39+
"description": "URL identifying the destination agent (for sales agents, etc.)"
40+
},
41+
"account": {
42+
"type": "string",
43+
"description": "Optional account identifier on the agent"
44+
}
45+
},
46+
"required": [
47+
"type",
48+
"agent_url"
49+
],
50+
"additionalProperties": false
51+
}
52+
]
53+
}

scripts/generate_models_simple.py

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,102 @@ def escape_string_for_python(text: str) -> str:
6666
return text.strip()
6767

6868

69+
def generate_discriminated_union(schema: dict, base_name: str) -> str:
70+
"""
71+
Generate Pydantic models for a discriminated union (oneOf with type discriminator).
72+
73+
Creates a base model for each variant and a union type for the parent.
74+
For example, Destination = PlatformDestination | AgentDestination
75+
"""
76+
lines = []
77+
78+
# Add schema description as a comment
79+
if "description" in schema:
80+
desc = escape_string_for_python(schema["description"])
81+
lines.append(f"# {desc}")
82+
lines.append("")
83+
84+
variant_names = []
85+
86+
# Generate a model for each variant in oneOf
87+
for i, variant in enumerate(schema.get("oneOf", [])):
88+
# Try to get discriminator value for better naming
89+
discriminator_value = None
90+
if "properties" in variant and "type" in variant["properties"]:
91+
type_prop = variant["properties"]["type"]
92+
if "const" in type_prop:
93+
discriminator_value = type_prop["const"]
94+
95+
# Generate variant name
96+
if discriminator_value:
97+
variant_name = f"{discriminator_value.capitalize()}{base_name}"
98+
else:
99+
variant_name = f"{base_name}Variant{i+1}"
100+
101+
variant_names.append(variant_name)
102+
103+
# Generate the variant model
104+
lines.append(f"class {variant_name}(BaseModel):")
105+
106+
# Add description if available
107+
if "description" in variant:
108+
desc = variant["description"].replace("\\", "\\\\").replace('"""', '\\"\\"\\"')
109+
desc = desc.replace("\n", " ").replace("\r", "")
110+
desc = re.sub(r"\s+", " ", desc).strip()
111+
lines.append(f' """{desc}"""')
112+
lines.append("")
113+
114+
# Add properties
115+
if "properties" in variant and variant["properties"]:
116+
for prop_name, prop_schema in variant["properties"].items():
117+
safe_name, needs_alias = sanitize_field_name(prop_name)
118+
prop_type = get_python_type(prop_schema)
119+
desc = prop_schema.get("description", "")
120+
if desc:
121+
desc = escape_string_for_python(desc)
122+
123+
is_required = prop_name in variant.get("required", [])
124+
125+
if is_required:
126+
if desc and needs_alias:
127+
lines.append(
128+
f' {safe_name}: {prop_type} = Field(alias="{prop_name}", description="{desc}")'
129+
)
130+
elif desc:
131+
lines.append(f' {safe_name}: {prop_type} = Field(description="{desc}")')
132+
elif needs_alias:
133+
lines.append(f' {safe_name}: {prop_type} = Field(alias="{prop_name}")')
134+
else:
135+
lines.append(f" {safe_name}: {prop_type}")
136+
else:
137+
if desc and needs_alias:
138+
lines.append(
139+
f' {safe_name}: {prop_type} | None = Field(None, alias="{prop_name}", description="{desc}")'
140+
)
141+
elif desc:
142+
lines.append(
143+
f' {safe_name}: {prop_type} | None = Field(None, description="{desc}")'
144+
)
145+
elif needs_alias:
146+
lines.append(
147+
f' {safe_name}: {prop_type} | None = Field(None, alias="{prop_name}")'
148+
)
149+
else:
150+
lines.append(f" {safe_name}: {prop_type} | None = None")
151+
else:
152+
lines.append(" pass")
153+
154+
lines.append("")
155+
lines.append("")
156+
157+
# Create union type
158+
union_type = " | ".join(variant_names)
159+
lines.append(f"# Union type for {schema.get('title', base_name)}")
160+
lines.append(f"{base_name} = {union_type}")
161+
162+
return "\n".join(lines)
163+
164+
69165
def generate_model_for_schema(schema_file: Path) -> str:
70166
"""Generate Pydantic model code for a single schema inline."""
71167
with open(schema_file) as f:
@@ -74,6 +170,10 @@ def generate_model_for_schema(schema_file: Path) -> str:
74170
# Start with model name
75171
model_name = snake_to_pascal(schema_file.stem)
76172

173+
# Check if this is a oneOf discriminated union
174+
if "oneOf" in schema and "properties" not in schema:
175+
return generate_discriminated_union(schema, model_name)
176+
77177
# Check if this is a simple type alias (enum or primitive type without properties)
78178
if "properties" not in schema:
79179
# This is a type alias, not a model class
@@ -155,6 +255,13 @@ def get_python_type(schema: dict) -> str:
155255
ref = schema["$ref"]
156256
return snake_to_pascal(ref.replace(".json", ""))
157257

258+
# Handle const (discriminator values)
259+
if "const" in schema:
260+
const_value = schema["const"]
261+
if isinstance(const_value, str):
262+
return f'Literal["{const_value}"]'
263+
return f"Literal[{const_value}]"
264+
158265
schema_type = schema.get("type")
159266

160267
if schema_type == "string":
@@ -387,6 +494,8 @@ def main():
387494
"protocol-envelope.json",
388495
"response.json",
389496
"promoted-products.json",
497+
"destination.json",
498+
"deployment.json",
390499
# Enum types (need type aliases)
391500
"channels.json",
392501
"delivery-type.json",
@@ -436,6 +545,7 @@ def main():
436545
"",
437546
"# These types are referenced in schemas but don't have schema files",
438547
"# Defining them as type aliases to maintain type safety",
548+
"ActivationKey = dict[str, Any]",
439549
"PackageRequest = dict[str, Any]",
440550
"PushNotificationConfig = dict[str, Any]",
441551
"ReportingCapabilities = dict[str, Any]",

0 commit comments

Comments
 (0)