3
3
import json
4
4
import sys
5
5
import time
6
- from typing import Any , Callable , ClassVar , Dict , List , Optional , Set , cast
6
+ from typing import Any , Callable , ClassVar , Dict , List , Optional , Set , Union , cast
7
7
8
8
from pydantic .v1 import Field , PrivateAttr
9
9
@@ -652,7 +652,7 @@ def applicable_entity_types(self, entity_types: EntityTypes):
652
652
self .options .applicable_entity_types = json .dumps (list (entity_types ))
653
653
654
654
@property
655
- def applicable_asset_types (self ) -> AssetTypes :
655
+ def applicable_asset_types (self ) -> Union [ Set [ str ], AssetTypes ] :
656
656
"""
657
657
Asset type names to which to restrict the attribute.
658
658
Only assets of one of these types will have this attribute available.
@@ -663,25 +663,27 @@ def applicable_asset_types(self) -> AssetTypes:
663
663
return set ()
664
664
665
665
@applicable_asset_types .setter
666
- def applicable_asset_types (self , asset_types ):
666
+ def applicable_asset_types (self , asset_types : Union [ Set [ str ], AssetTypes ] ):
667
667
if self .options is None :
668
668
raise ErrorCode .MISSING_OPTIONS .exception_with_parameters ()
669
+
669
670
if not isinstance (asset_types , set ):
670
671
raise ErrorCode .INVALID_PARAMETER_TYPE .exception_with_parameters (
671
672
"applicable_asset_types" , AssetTypes
672
673
)
673
- invalid_types = set ()
674
- for asset_type in asset_types :
674
+
675
+ # Validate asset type names against SDK asset classes
676
+ invalid_types = {
677
+ asset_type
678
+ for asset_type in asset_types
675
679
if not getattr (
676
680
sys .modules .get ("pyatlan.model.assets" , {}), asset_type , None
677
- ):
678
- invalid_types .add (asset_type )
679
-
681
+ )
682
+ }
680
683
if invalid_types :
681
684
raise ErrorCode .INVALID_PARAMETER_VALUE .exception_with_parameters (
682
- asset_types , "applicable_asset_types" , "SDK asset types"
685
+ invalid_types , "applicable_asset_types" , "SDK asset types"
683
686
)
684
-
685
687
self .options .applicable_asset_types = json .dumps (list (asset_types ))
686
688
687
689
@property
@@ -835,7 +837,7 @@ def create(
835
837
multi_valued : bool = False ,
836
838
options_name : Optional [str ] = None ,
837
839
applicable_connections : Optional [Set [str ]] = None ,
838
- applicable_asset_types : Optional [AssetTypes ] = None ,
840
+ applicable_asset_types : Optional [Union [ Set [ str ], AssetTypes ] ] = None ,
839
841
applicable_glossaries : Optional [Set [str ]] = None ,
840
842
applicable_glossary_types : Optional [GlossaryTypes ] = None ,
841
843
applicable_other_asset_types : Optional [OtherAssetTypes ] = None ,
0 commit comments