3030_failed_imports : set [str ] = set ()
3131
3232
33- def _get_type_name_default (cls : type ) -> str | None :
34- """Extract the type_name default value from a msgspec Struct class.
35-
36- For msgspec Structs, class attributes become field descriptors, so we need
37- to use msgspec.structs.fields() to get the actual default value.
38- """
39- try :
40- for field in msgspec .structs .fields (cls ):
41- if field .name == "type_name" :
42- default = field .default
43- if isinstance (default , str ) and default != "UNSET" :
44- return default
45- return None
46- except TypeError :
47- # Not a msgspec Struct - fall back to getattr
48- pass
49-
50- # Fallback for non-Struct classes
51- type_name = getattr (cls , "type_name" , None )
52- if isinstance (type_name , str ) and type_name != "UNSET" :
53- return type_name
54- return None
55-
56-
5733def register_asset (cls : _T ) -> _T :
5834 """Decorator that registers an Asset subclass in the type registry.
5935
@@ -65,15 +41,12 @@ def register_asset(cls: _T) -> _T:
6541
6642 @register_asset
6743 class AtlasGlossaryTerm(Asset):
68- type_name: Union[str, UnsetType] = "AtlasGlossaryTerm"
6944 ...
7045
71- The class is registered under its `type_name` field's default value.
72- Works with both msgspec Structs and regular classes .
46+ The class is registered under its class name, which matches the Atlas
47+ typeName (e.g., class Table -> typeName "Table") .
7348 """
74- type_name = _get_type_name_default (cls )
75- if type_name :
76- _type_registry [type_name ] = cls
49+ _type_registry [cls .__name__ ] = cls
7750 return cls
7851
7952
@@ -244,25 +217,6 @@ def to_atlas_format(asset: Asset) -> dict[str, Any]:
244217 ("attributes" , "uniqueAttributes" , "relationshipAttributes" )
245218)
246219
247- _CAMEL_ABBREV_RE = re .compile (r"([A-Z]{2,})(?=[A-Z][a-z]|$)" )
248-
249- # Keys that use explicit msgspec.field(name="...") with trailing uppercase
250- # (e.g. dataProductAssetsDSL). Normalization would turn them into Dsl;
251- # preserve so struct field name matches.
252- _PRESERVE_CAMEL_KEYS = frozenset ({"dataProductAssetsDSL" , "apiPathRawURI" })
253-
254-
255- def _normalize_camel_key (key : str ) -> str :
256- """Normalize uppercase abbreviations in camelCase keys for msgspec.
257-
258- msgspec's rename="camel" expects apiPathRawUri, not apiPathRawURI.
259- This converts trailing/mid uppercase runs like URI→Uri, DSL→Dsl, DQ→Dq.
260- Keys in _PRESERVE_CAMEL_KEYS are left unchanged so they match struct fields.
261- """
262- if key in _PRESERVE_CAMEL_KEYS :
263- return key
264- return _CAMEL_ABBREV_RE .sub (lambda m : m .group (1 ).capitalize (), key )
265-
266220
267221def _flatten_entity_dict (data : dict [str , Any ]) -> dict [str , Any ]:
268222 """Flatten one Atlas entity dict, merging ``attributes``,
@@ -276,9 +230,9 @@ def _flatten_entity_dict(data: dict[str, Any]) -> dict[str, Any]:
276230 if key in _NESTED_BUCKETS :
277231 if isinstance (value , dict ):
278232 for k , v in value .items ():
279- flattened [_normalize_camel_key ( k ) ] = v
233+ flattened [k ] = v
280234 else :
281- flattened [_normalize_camel_key ( key ) ] = value
235+ flattened [key ] = value
282236
283237 for key , value in list (flattened .items ()):
284238 if isinstance (value , dict ) and "typeName" in value :
0 commit comments