1
- from labelbox .utils import camel_case
2
- from pydantic import BaseModel
1
+ from dataclasses import dataclass
3
2
3
+ from labelbox .utils import snake_case
4
4
from labelbox .orm .db_object import DbObject
5
5
from labelbox .orm .model import Field
6
6
7
7
8
-
9
-
10
-
11
- class AwsIamIntegrationSettings (BaseModel ):
8
+ @dataclass
9
+ class AwsIamIntegrationSettings :
12
10
role_arn : str
13
11
14
- class Config :
15
- allow_population_by_field_name = True
16
- alias_generator = camel_case
17
-
18
12
19
- class GcpIamIntegrationSettings (BaseModel ):
13
+ @dataclass
14
+ class GcpIamIntegrationSettings :
20
15
service_account_email_id : str
21
16
read_bucket : str
22
17
23
- class Config :
24
- allow_population_by_field_name = True
25
- alias_generator = camel_case
26
-
27
-
28
18
29
19
class IAMIntegration (DbObject ):
30
20
""" Represents an IAM integration for delegated access
@@ -41,13 +31,18 @@ class IAMIntegration(DbObject):
41
31
"""
42
32
43
33
def __init__ (self , client , data ):
44
- settings = data .pop ('settings' , {})
45
- type_name = settings .pop ('__typename' )
46
- if type_name == "GcpIamIntegrationSettings" :
47
- self .settings = GcpIamIntegrationSettings (** settings )
48
- elif type_name == "AwsIamIntegrationSettings" :
49
- self .settings = AwsIamIntegrationSettings (** settings )
50
-
34
+ settings = data .pop ('settings' , None )
35
+ if settings is not None :
36
+ type_name = settings .pop ('__typename' )
37
+ settings = {snake_case (k ): v for k , v in settings .items ()}
38
+ if type_name == "GcpIamIntegrationSettings" :
39
+ self .settings = GcpIamIntegrationSettings (** settings )
40
+ elif type_name == "AwsIamIntegrationSettings" :
41
+ self .settings = AwsIamIntegrationSettings (** settings )
42
+ else :
43
+ self .settings = None
44
+ else :
45
+ self .settings = None
51
46
super ().__init__ (client , data )
52
47
53
48
_DEFAULT = "DEFAULT"
0 commit comments