Coverage for trimesh/exchange/gltf/extensions.py: 70%
63 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-24 04:40 +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
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 handle_extensions(
158 *,
159 extensions: dict[str, Any] | None,
160 scope: Scope,
161 **kwargs,
162) -> Any:
163 """
164 Process extensions dict for a given scope, calling registered handlers.
166 Parameters
167 ----------
168 extensions
169 The "extensions" dict from a glTF element, or None.
170 scope
171 Handler scope to invoke.
172 **kwargs
173 Scope-specific arguments that will be combined with extension data
174 into a typed context dict. Required kwargs by scope:
175 - material: parse_textures, images
176 - texture_source: (none)
177 - primitive: primitive, mesh_kwargs, accessors
178 - primitive_preprocess: primitive, accessors, views
179 - primitive_export: mesh, name, tree, buffer_items, primitive, include_normals
181 Returns
182 -------
183 results
184 Dict of {extension_name: result} for most scopes.
185 For scopes ending in "_source", returns first non-None result.
186 For "primitive" scope, automatically merges results into mesh_kwargs.
187 """
188 if not extensions or scope not in _handlers:
189 return {} if not scope.endswith("_source") else None
191 results = {}
192 for ext_name, data in extensions.items():
193 if ext_name not in _handlers[scope]:
194 continue
195 try:
196 # Build context dict with data + all kwargs
197 context = {"data": data, **kwargs}
198 if (result := _handlers[scope][ext_name](context)) is not None:
199 results[ext_name] = result
200 except Exception as e:
201 log.warning(f"failed to process extension {ext_name}: {e}")
203 # for _source scopes return first result, otherwise return all results
204 if scope.endswith("_source"):
205 return next(iter(results.values()), None)
207 # for primitive scope, automatically merge results into mesh_kwargs
208 if scope == "primitive" and "mesh_kwargs" in kwargs:
209 mesh_kwargs = kwargs["mesh_kwargs"]
210 for ext_result in results.values():
211 if not isinstance(ext_result, dict):
212 continue
213 # merge extension results, recursively merging nested dicts
214 for key, value in ext_result.items():
215 if isinstance(value, dict):
216 if key not in mesh_kwargs:
217 mesh_kwargs[key] = {}
218 _deep_merge(mesh_kwargs[key], value)
219 else:
220 mesh_kwargs[key] = value
222 return results
225# ----------------------------------------------------------------------
226# Built-in handlers
227# ----------------------------------------------------------------------
230@register_handler("KHR_materials_pbrSpecularGlossiness", scope="material")
231def _specular_glossiness(context: MaterialContext) -> dict[str, Any] | None:
232 """
233 Convert specular-glossiness material to PBR metallic-roughness.
235 Parameters
236 ----------
237 context
238 MaterialContext with extension data, parse_textures function, and images.
240 Returns
241 -------
242 pbr_dict
243 PBR metallic-roughness parameters, or None on failure.
244 """
245 try:
246 from ...visual.gloss import specular_to_pbr
248 return specular_to_pbr(**context["parse_textures"](data=context["data"]))
249 except Exception:
250 log.debug("failed to convert specular-glossiness", exc_info=True)
251 return None
254@register_handler("EXT_texture_webp", scope="texture_source")
255def _texture_webp_source(context: TextureSourceContext) -> int | None:
256 """
257 Return image source index from EXT_texture_webp.
259 Parameters
260 ----------
261 context
262 TextureSourceContext with extension data.
264 Returns
265 -------
266 source_index
267 Index into glTF images array, or None if not present.
268 """
269 return context["data"].get("source")