2323logger .info ("Generating schema reference..." )
2424
2525
26- def get_type (annotation : Type ) -> str :
26+ def _is_linkable_type (annotation : Any ) -> bool :
27+ """Check if a type annotation contains a BaseModel subclass (excluding Range)."""
28+ if inspect .isclass (annotation ):
29+ return issubclass (annotation , BaseModel ) and not issubclass (annotation , Range )
30+ origin = get_origin (annotation )
31+ if origin is Annotated :
32+ return _is_linkable_type (get_args (annotation )[0 ])
33+ if origin is Union :
34+ return any (_is_linkable_type (arg ) for arg in get_args (annotation ))
35+ if origin is list :
36+ args = get_args (annotation )
37+ return bool (args ) and _is_linkable_type (args [0 ])
38+ return False
39+
40+
41+ def _type_sort_key (t : str ) -> tuple :
42+ """Sort key for type parts: primitives first, then literals, then compound types."""
43+ order = {"bool" : 0 , "int" : 1 , "float" : 2 , "str" : 3 }
44+ if t in order :
45+ return (0 , order [t ])
46+ if t .startswith ('"' ):
47+ return (1 , t )
48+ if t .startswith ("list" ):
49+ return (2 , t )
50+ if t == "dict" :
51+ return (3 , "" )
52+ if t == "object" :
53+ return (4 , "" )
54+ return (5 , t )
55+
56+
57+ def get_friendly_type (annotation : Type ) -> str :
58+ """Get a user-friendly type string for documentation.
59+
60+ Produces types like: ``int | str``, ``"rps"``, ``list[object]``, ``"spot" | "on-demand" | "auto"``.
61+ """
62+ # Unwrap Annotated
2763 if get_origin (annotation ) is Annotated :
28- return get_type (get_args (annotation )[0 ])
64+ return get_friendly_type (get_args (annotation )[0 ])
65+
66+ # Handle Union (including Optional)
2967 if get_origin (annotation ) is Union :
30- # Optional is Union with None.
31- # We don't want to show Optional[A, None] but just Optional[A]
32- if annotation .__name__ == "Optional" :
33- args = "," .join (get_type (arg ) for arg in get_args (annotation )[:- 1 ])
34- else :
35- args = "," .join (get_type (arg ) for arg in get_args (annotation ))
36- return f"{ annotation .__name__ } { args }
68+ args = [a for a in get_args (annotation ) if a is not type (None )]
69+ if not args :
70+ return ""
71+ parts : list = []
72+ for arg in args :
73+ friendly = get_friendly_type (arg )
74+ # Split compound types (e.g., "int | str" from Range) to deduplicate,
75+ # but avoid splitting types that contain brackets (e.g., list[...])
76+ if "[" not in friendly :
77+ for part in friendly .split (" | " ):
78+ if part and part not in parts :
79+ parts .append (part )
80+ else :
81+ if friendly and friendly not in parts :
82+ parts .append (friendly )
83+ parts .sort (key = _type_sort_key )
84+ return " | " .join (parts )
85+
86+ # Handle Literal — show as enum (specific values are in the field description)
3787 if get_origin (annotation ) is Literal :
38- return str (annotation ).split ("." , maxsplit = 1 )[- 1 ]
88+ return "enum"
89+
90+ # Handle list
3991 if get_origin (annotation ) is list :
40- return f"List[{ get_type (get_args (annotation )[0 ])}
92+ args = get_args (annotation )
93+ if args :
94+ inner = get_friendly_type (args [0 ])
95+ return f"list[{ inner }
96+ return "list"
97+
98+ # Handle dict
4199 if get_origin (annotation ) is dict :
42- return f"Dict[{ get_type (get_args (annotation )[0 ])} { get_type (get_args (annotation )[1 ])}
43- return annotation .__name__
100+ return "dict"
101+
102+ # Handle concrete classes
103+ if inspect .isclass (annotation ):
104+ # Enum — list values
105+ if issubclass (annotation , Enum ):
106+ values = [e .value for e in annotation ]
107+ return " | " .join (f'"{ v } for v in values )
108+
109+ # Range — depends on inner type parameter
110+ if issubclass (annotation , Range ):
111+ min_field = annotation .__fields__ .get ("min" )
112+ if min_field and inspect .isclass (min_field .type_ ):
113+ # Range[Memory] → str, Range[int] → int | str
114+ if issubclass (min_field .type_ , float ):
115+ return "str"
116+ return "int | str"
117+
118+ # Memory (float subclass that parses "8GB" strings)
119+ from dstack ._internal .core .models .resources import Memory as _Memory
120+
121+ if issubclass (annotation , _Memory ):
122+ return "str"
123+
124+ # BaseModel subclass (not Range)
125+ if issubclass (annotation , BaseModel ) and not issubclass (annotation , Range ):
126+ # Root models (with __root__ field) — resolve from the root type
127+ if "__root__" in annotation .__fields__ :
128+ return get_friendly_type (annotation .__fields__ ["__root__" ].annotation )
129+ # Models with custom __get_validators__ accept primitive input (int, str)
130+ # in addition to the full object form (e.g., GPUSpec, CPUSpec, DiskSpec)
131+ if "__get_validators__" in annotation .__dict__ :
132+ return "int | str | object"
133+ return "object"
134+
135+ # ComputeCapability (tuple subclass that parses "7.5" strings)
136+ if annotation .__name__ == "ComputeCapability" :
137+ return "float | str"
138+
139+ # Constrained and primitive types — check MRO
140+ # bool must come before int (bool is a subclass of int)
141+ if issubclass (annotation , bool ):
142+ return "bool"
143+ if issubclass (annotation , int ):
144+ # Duration (int subclass that parses "5m" strings)
145+ if annotation .__name__ == "Duration" :
146+ return "int | str"
147+ return "int"
148+ if issubclass (annotation , float ):
149+ return "float"
150+ if issubclass (annotation , str ):
151+ return "str"
152+ if issubclass (annotation , (list , tuple )):
153+ return "list"
154+ if issubclass (annotation , dict ):
155+ return "dict"
156+
157+ return annotation .__name__
158+
159+ return str (annotation )
160+
161+
162+ _JSON_SCHEMA_TYPE_MAP = {
163+ "string" : "str" ,
164+ "integer" : "int" ,
165+ "number" : "float" ,
166+ "boolean" : "bool" ,
167+ "array" : "list" ,
168+ "object" : "object" ,
169+ }
170+
171+
172+ def _enrich_type_from_schema (friendly_type : str , prop_schema : Dict [str , Any ]) -> str :
173+ """Enrich the friendly type with extra accepted types from the JSON schema.
174+
175+ Models may define ``schema_extra`` that adds ``anyOf`` entries for fields
176+ that accept alternative input types (e.g., duration fields typed as ``int``
177+ but also accepting ``str`` like ``"5m"``).
178+ """
179+ any_of = prop_schema .get ("anyOf" )
180+ if not any_of :
181+ return friendly_type
182+ # Only consider string/integer — the most common alternative input types.
183+ # Skip boolean (typically a backward-compat artifact) and object/array.
184+ _ENRICHABLE = {"string" : "str" , "integer" : "int" }
185+ schema_types = set ()
186+ for entry in any_of :
187+ mapped = _ENRICHABLE .get (entry .get ("type" , "" ))
188+ if mapped :
189+ schema_types .add (mapped )
190+ # Add any schema types not already present in the friendly type
191+ current_parts = [p .strip () for p in friendly_type .split (" | " )]
192+ new_parts = schema_types - set (current_parts )
193+ if not new_parts :
194+ return friendly_type
195+ all_parts = list (set (current_parts ) | new_parts )
196+ # If str is now present, enum is redundant
197+ if "str" in all_parts and "enum" in all_parts :
198+ all_parts .remove ("enum" )
199+ all_parts .sort (key = _type_sort_key )
200+ return " | " .join (all_parts )
44201
45202
46203def generate_schema_reference (
@@ -63,14 +220,21 @@ def generate_schema_reference(
63220 "" ,
64221 ]
65222 )
223+ # Get JSON schema to detect extra accepted types from schema_extra
224+ try :
225+ schema_props = cls .schema ().get ("properties" , {})
226+ except Exception :
227+ schema_props = {}
66228 for name , field in cls .__fields__ .items ():
67229 default = field .default
68230 if isinstance (default , Enum ):
69231 default = default .value
232+ friendly_type = get_friendly_type (field .annotation )
233+ friendly_type = _enrich_type_from_schema (friendly_type , schema_props .get (name , {}))
70234 values = dict (
71235 name = name ,
72236 description = field .field_info .description ,
73- type = get_type ( field . annotation ) ,
237+ type = friendly_type ,
74238 default = default ,
75239 required = field .required ,
76240 )
@@ -84,11 +248,7 @@ def generate_schema_reference(
84248 if field .annotation .__name__ == "Annotated" :
85249 if field_type .__name__ in ["Optional" , "List" , "list" , "Union" ]:
86250 field_type = get_args (field_type )[0 ]
87- base_model = (
88- inspect .isclass (field_type )
89- and issubclass (field_type , BaseModel )
90- and not issubclass (field_type , Range )
91- )
251+ base_model = _is_linkable_type (field_type )
92252 else :
93253 base_model = False
94254 _defaults = (
@@ -114,29 +274,27 @@ def generate_schema_reference(
114274 if not base_model
115275 else f"[`{ values ['name' ]} { item_id_prefix } { link_name }
116276 )
117- item_optional_marker = "(Optional)" if not values ["required" ] else ""
277+ item_required_marker = "(Required)" if values ["required" ] else "(Optional)"
278+ item_type_display = f"`{ values ['type' ]} if values .get ("type" ) else ""
118279 item_description = (values ["description" ]).replace ("\n " , "<br>" ) + "."
119280 item_default = _defaults if not values ["required" ] else _must_be
120281 item_id = f"#{ values ['name' ]} if not base_model else f"#_{ values ['name' ]}
121282 item_toc_label = f"data-toc-label='{ values ['name' ]}
122283 item_css_cass = "class='reference-item'"
123- rows .append (
124- prefix
125- + " " .join (
126- [
127- f"###### { item_header } ,
128- "-" ,
129- item_optional_marker ,
130- item_description ,
131- item_default ,
132- "{" ,
133- item_id ,
134- item_toc_label ,
135- item_css_cass ,
136- "}" ,
137- ]
138- )
139- )
284+ parts = [
285+ f"###### { item_header } ,
286+ "-" ,
287+ item_required_marker ,
288+ item_type_display ,
289+ item_description ,
290+ item_default ,
291+ "{" ,
292+ item_id ,
293+ item_toc_label ,
294+ item_css_cass ,
295+ "}" ,
296+ ]
297+ rows .append (prefix + " " .join (p for p in parts if p ))
140298 return "\n " .join (rows )
141299
142300
0 commit comments