Skip to content

Commit d4e7945

Browse files
Aryamanz29claude
andcommitted
BLDX-709, BLDX-710 | Clean up transform.py
- Remove _normalize_camel_key, _PRESERVE_CAMEL_KEYS, _CAMEL_ABBREV_RE (msgspec.field(name="...") on generated structs handles these cases) - Simplify register_asset to use cls.__name__ instead of reading field defaults (supports BLDX-710 where type_name default becomes UNSET) Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent bd44034 commit d4e7945

1 file changed

Lines changed: 5 additions & 51 deletions

File tree

pyatlan_v9/model/transform.py

Lines changed: 5 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -30,30 +30,6 @@
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-
5733
def 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

267221
def _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

Comments
 (0)