Coverage for trimesh/exchange/gltf/extensions.py: 77%
65 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-07-31 23:55 +0000
1"""
2gltf_extensions.py
3------------------
5Extension registry for glTF import/export with scope-based handlers.
6Each scope has a TypedDict defining the context passed to handlers.
7"""
9from collections import OrderedDict
10from collections.abc import Callable, Iterable
11from typing import Any, Literal, TypeAlias, TypedDict
13from ...constants import log
15# Scopes define where in the glTF load/export process handlers run:
16# material - after parsing material, can override PBR values
17# texture_source - when resolving texture image index
18# primitive - after loading primitive, can add face_attributes
19# primitive_preprocess - before accessor reads, can modify accessors in-place
20# primitive_export - during mesh export, can compress/modify primitive data
21Scope: TypeAlias = Literal[
22 "material", "texture_source", "primitive", "primitive_preprocess", "primitive_export"
23]
26# ----------------------------------------------------------------------
27# TypedDict contexts for each scope
28# ----------------------------------------------------------------------
29#
30# These TypedDicts define the MINIMUM fields passed to handlers for each scope.
31# Additional fields may be added in future versions for new functionality.
32#
33# FOR FORWARD COMPATIBILITY: Handlers should access only the fields they need
34# and ignore unknown fields. The context is passed as a plain dict at runtime,
35# so handlers can safely use dict.get() for optional access or simply not
36# reference fields they don't need.
37#
38# Example handler pattern:
39#
40# def my_handler(context: MaterialContext) -> dict | None:
41# # Access only what you need - additional fields won't break this
42# data = context["data"]
43# images = context["images"]
44# return {"baseColorFactor": [1, 0, 0, 1]}
45#
46# ----------------------------------------------------------------------
49class MaterialContext(TypedDict):
50 """Context for material scope handlers."""
52 data: dict[str, Any]
53 parse_textures: Callable[..., dict[str, Any]]
54 images: list
57class TextureSourceContext(TypedDict):
58 """Context for texture_source scope handlers."""
60 data: dict[str, Any]
63class PrimitiveContext(TypedDict):
64 """Context for primitive scope handlers (post-load)."""
66 data: dict[str, Any]
67 primitive: dict
68 mesh_kwargs: dict
69 accessors: list
72class PrimitivePreprocessContext(TypedDict):
73 """Context for primitive_preprocess scope handlers (pre-load)."""
75 data: dict[str, Any]
76 primitive: dict
77 accessors: list
78 views: list
81class PrimitiveExportContext(TypedDict):
82 """Context for primitive_export scope handlers (during export)."""
84 mesh: Any
85 name: str
86 tree: dict
87 buffer_items: OrderedDict
88 primitive: dict
89 include_normals: bool
92# Handler type alias - handlers receive a context dict
93Handler: TypeAlias = Callable[[Any], Any]
95# callback to parse material dict and resolve texture references
96# signature: (*, data: dict) -> dict
97ParseTextures: TypeAlias = Callable[..., dict[str, Any]]
99# Registry: {scope: {extension_name: handler}}
100_handlers: dict[str, dict[str, Handler]] = {}
103def _deep_merge(target: dict, source: dict) -> None:
104 """
105 Recursively merge source dict into target dict.
107 Parameters
108 ----------
109 target
110 Dict to merge into (modified in place)
111 source
112 Dict to merge from
113 """
114 for key, value in source.items():
115 if isinstance(value, dict) and key in target and isinstance(target[key], dict):
116 # Both are dicts - recurse
117 _deep_merge(target[key], value)
118 else:
119 # Overwrite or set new key
120 target[key] = value
123def register_handler(name: str, scope: Scope) -> Callable[[Handler], Handler]:
124 """
125 Decorator to register a handler for a glTF extension.
127 Parameters
128 ----------
129 name
130 Extension name, e.g. "KHR_materials_pbrSpecularGlossiness".
131 scope
132 Handler scope, e.g. "material", "texture_source", "primitive".
134 Returns
135 -------
136 decorator
137 Function that registers the handler and returns it unchanged.
139 Example
140 -------
141 >>> @register_handler("MY_extension", scope="material")
142 ... def my_handler(context: MaterialContext) -> dict | None:
143 ... data = context["data"]
144 ... images = context["images"]
145 ... return {"baseColorFactor": [1, 0, 0, 1]}
146 """
147 if scope not in _handlers:
148 _handlers[scope] = {}
150 def decorator(func: Handler) -> Handler:
151 _handlers[scope][name] = func
152 return func
154 return decorator
157def unregistered(extensions: Iterable[str], scope: Scope) -> set:
158 """
159 Find extension names with no registered handler for a scope.
161 Parameters
162 ----------
163 extensions
164 Extension names, i.e. the keys of a glTF "extensions" dict.
165 scope
166 Handler scope to check against.
168 Returns
169 -------
170 missing
171 Extension names with no handler registered for the scope.
172 """
173 return set(extensions) - _handlers.get(scope, {}).keys()
176def handle_extensions(
177 *,
178 extensions: dict[str, Any] | None,
179 scope: Scope,
180 **kwargs,
181) -> Any:
182 """
183 Process extensions dict for a given scope, calling registered handlers.
185 Parameters
186 ----------
187 extensions
188 The "extensions" dict from a glTF element, or None.
189 scope
190 Handler scope to invoke.
191 **kwargs
192 Scope-specific arguments that will be combined with extension data
193 into a typed context dict. Required kwargs by scope:
194 - material: parse_textures, images
195 - texture_source: (none)
196 - primitive: primitive, mesh_kwargs, accessors
197 - primitive_preprocess: primitive, accessors, views
198 - primitive_export: mesh, name, tree, buffer_items, primitive, include_normals
200 Returns
201 -------
202 results
203 Dict of {extension_name: result} for most scopes.
204 For scopes ending in "_source", returns first non-None result.
205 For "primitive" scope, automatically merges results into mesh_kwargs.
206 """
207 if not extensions or scope not in _handlers:
208 return {} if not scope.endswith("_source") else None
210 results = {}
211 for ext_name, data in extensions.items():
212 if ext_name not in _handlers[scope]:
213 continue
214 try:
215 # Build context dict with data + all kwargs
216 context = {"data": data, **kwargs}
217 if (result := _handlers[scope][ext_name](context)) is not None:
218 results[ext_name] = result
219 except Exception as e:
220 log.warning(f"failed to process extension {ext_name}: {e}")
222 # for _source scopes return first result, otherwise return all results
223 if scope.endswith("_source"):
224 return next(iter(results.values()), None)
226 # for primitive scope, automatically merge results into mesh_kwargs
227 if scope == "primitive" and "mesh_kwargs" in kwargs:
228 mesh_kwargs = kwargs["mesh_kwargs"]
229 for ext_result in results.values():
230 if not isinstance(ext_result, dict):
231 continue
232 # merge extension results, recursively merging nested dicts
233 for key, value in ext_result.items():
234 if isinstance(value, dict):
235 if key not in mesh_kwargs:
236 mesh_kwargs[key] = {}
237 _deep_merge(mesh_kwargs[key], value)
238 else:
239 mesh_kwargs[key] = value
241 return results
244# ----------------------------------------------------------------------
245# Built-in handlers
246# ----------------------------------------------------------------------
249@register_handler("KHR_materials_pbrSpecularGlossiness", scope="material")
250def _specular_glossiness(context: MaterialContext) -> dict[str, Any] | None:
251 """
252 Convert specular-glossiness material to PBR metallic-roughness.
254 Parameters
255 ----------
256 context
257 MaterialContext with extension data, parse_textures function, and images.
259 Returns
260 -------
261 pbr_dict
262 PBR metallic-roughness parameters, or None on failure.
263 """
264 try:
265 from ...visual.gloss import specular_to_pbr
267 return specular_to_pbr(**context["parse_textures"](data=context["data"]))
268 except Exception:
269 log.debug("failed to convert specular-glossiness", exc_info=True)
270 return None
273@register_handler("EXT_texture_webp", scope="texture_source")
274def _texture_webp_source(context: TextureSourceContext) -> int | None:
275 """
276 Return image source index from EXT_texture_webp.
278 Parameters
279 ----------
280 context
281 TextureSourceContext with extension data.
283 Returns
284 -------
285 source_index
286 Index into glTF images array, or None if not present.
287 """
288 return context["data"].get("source")