@@ -41,7 +41,7 @@ class AggregationType(str, enum.Enum):
41
41
join_existing = 'join_existing'
42
42
union = 'union'
43
43
44
- model_config = ConfigDict (validate_default = True , validate_assignment = True )
44
+ model_config = ConfigDict (validate_assignment = True )
45
45
46
46
47
47
class DataFormat (str , enum .Enum ):
@@ -50,22 +50,22 @@ class DataFormat(str, enum.Enum):
50
50
reference = 'reference'
51
51
opendap = 'opendap'
52
52
53
- model_config = ConfigDict (validate_default = True , validate_assignment = True )
53
+ model_config = ConfigDict (validate_assignment = True )
54
54
55
55
56
56
class Attribute (pydantic .BaseModel ):
57
57
column_name : pydantic .StrictStr
58
58
vocabulary : pydantic .StrictStr = ''
59
59
60
- model_config = ConfigDict (validate_default = True , validate_assignment = True )
60
+ model_config = ConfigDict (validate_assignment = True )
61
61
62
62
63
63
class Assets (pydantic .BaseModel ):
64
64
column_name : pydantic .StrictStr
65
- format : typing . Optional [ DataFormat ] = None
66
- format_column_name : typing . Optional [ pydantic .StrictStr ] = None
65
+ format : DataFormat | None = None
66
+ format_column_name : pydantic .StrictStr | None = None
67
67
68
- model_config = ConfigDict (validate_default = True , validate_assignment = True )
68
+ model_config = ConfigDict (validate_assignment = True )
69
69
70
70
@pydantic .model_validator (mode = 'after' )
71
71
def _validate_data_format (cls , model ):
@@ -82,7 +82,7 @@ class Aggregation(pydantic.BaseModel):
82
82
attribute_name : pydantic .StrictStr
83
83
options : dict = {}
84
84
85
- model_config = ConfigDict (validate_default = True , validate_assignment = True )
85
+ model_config = ConfigDict (validate_assignment = True )
86
86
87
87
88
88
class AggregationControl (pydantic .BaseModel ):
@@ -101,18 +101,16 @@ class ESMCatalogModel(pydantic.BaseModel):
101
101
esmcat_version : pydantic .StrictStr
102
102
attributes : list [Attribute ]
103
103
assets : Assets
104
- aggregation_control : typing . Optional [ AggregationControl ] = None
104
+ aggregation_control : AggregationControl | None = None
105
105
id : str = ''
106
- catalog_dict : typing . Optional [ list [dict ]] = None
107
- catalog_file : typing . Optional [ pydantic .StrictStr ] = None
108
- description : typing . Optional [ pydantic .StrictStr ] = None
109
- title : typing . Optional [ pydantic .StrictStr ] = None
110
- last_updated : typing . Optional [ typing . Union [ datetime .datetime , datetime .date ]] = None
106
+ catalog_dict : list [dict ] | None = None
107
+ catalog_file : pydantic .StrictStr | None = None
108
+ description : pydantic .StrictStr | None = None
109
+ title : pydantic .StrictStr | None = None
110
+ last_updated : datetime .datetime | datetime .date | None = None
111
111
_df : pd .DataFrame = pydantic .PrivateAttr ()
112
112
113
- model_config = ConfigDict (
114
- arbitrary_types_allowed = True , validate_default = True , validate_assignment = True
115
- )
113
+ model_config = ConfigDict (arbitrary_types_allowed = True , validate_assignment = True )
116
114
117
115
@pydantic .model_validator (mode = 'after' )
118
116
def validate_catalog (cls , model ):
@@ -136,11 +134,11 @@ def save(
136
134
self ,
137
135
name : str ,
138
136
* ,
139
- directory : str = None ,
137
+ directory : str | None = None ,
140
138
catalog_type : str = 'dict' ,
141
- to_csv_kwargs : dict = None ,
142
- json_dump_kwargs : dict = None ,
143
- storage_options : dict [str , typing .Any ] = None ,
139
+ to_csv_kwargs : dict | None = None ,
140
+ json_dump_kwargs : dict | None = None ,
141
+ storage_options : dict [str , typing .Any ] | None = None ,
144
142
) -> None :
145
143
"""
146
144
Save the catalog to a file.
@@ -193,7 +191,7 @@ def save(
193
191
194
192
if catalog_type == 'file' :
195
193
csv_kwargs = {'index' : False }
196
- csv_kwargs . update ( to_csv_kwargs or {})
194
+ csv_kwargs |= to_csv_kwargs or {}
197
195
compression = csv_kwargs .get ('compression' )
198
196
extensions = {'gzip' : '.gz' , 'bz2' : '.bz2' , 'zip' : '.zip' , 'xz' : '.xz' , None : '' }
199
197
csv_file_name = f'{ csv_file_name } { extensions [compression ]} '
@@ -206,15 +204,15 @@ def save(
206
204
207
205
with fs .open (json_file_name , 'w' ) as outfile :
208
206
json_kwargs = {'indent' : 2 }
209
- json_kwargs . update ( json_dump_kwargs or {})
207
+ json_kwargs |= json_dump_kwargs or {}
210
208
json .dump (data , outfile , ** json_kwargs )
211
209
212
210
print (f'Successfully wrote ESM catalog json file to: { json_file_name } ' )
213
211
214
212
@classmethod
215
213
def load (
216
214
cls ,
217
- json_file : typing . Union [ str , pydantic .FilePath , pydantic .AnyUrl ] ,
215
+ json_file : str | pydantic .FilePath | pydantic .AnyUrl ,
218
216
storage_options : dict [str , typing .Any ] = None ,
219
217
read_csv_kwargs : dict [str , typing .Any ] = None ,
220
218
) -> 'ESMCatalogModel' :
@@ -287,16 +285,20 @@ def _cast_agg_columns_with_iterables(self) -> None:
287
285
to avoid hashing issues (e.g. TypeError: unhashable type: 'list')
288
286
"""
289
287
if self .aggregation_control :
290
- columns = list (
288
+ if columns : = list (
291
289
self .columns_with_iterables .intersection (
292
- set (map (lambda agg : agg .attribute_name , self .aggregation_control .aggregations ))
290
+ set (
291
+ map (
292
+ lambda agg : agg .attribute_name ,
293
+ self .aggregation_control .aggregations ,
294
+ )
295
+ )
293
296
)
294
- )
295
- if columns :
297
+ ):
296
298
self ._df [columns ] = self ._df [columns ].apply (tuple )
297
299
298
300
@property
299
- def grouped (self ) -> typing . Union [ pd .core .groupby .DataFrameGroupBy , pd .DataFrame ] :
301
+ def grouped (self ) -> pd .core .groupby .DataFrameGroupBy | pd .DataFrame :
300
302
if self .aggregation_control :
301
303
if self .aggregation_control .groupby_attrs :
302
304
self .aggregation_control .groupby_attrs = list (
@@ -318,7 +320,7 @@ def grouped(self) -> typing.Union[pd.core.groupby.DataFrameGroupBy, pd.DataFrame
318
320
)
319
321
return self .df .groupby (cols )
320
322
321
- def _construct_group_keys (self , sep : str = '.' ) -> dict [str , typing . Union [ str , tuple [str ] ]]:
323
+ def _construct_group_keys (self , sep : str = '.' ) -> dict [str , str | tuple [str ]]:
322
324
internal_keys = self .grouped .groups .keys ()
323
325
public_keys = map (
324
326
lambda key : key if isinstance (key , str ) else sep .join (str (value ) for value in key ),
@@ -352,7 +354,7 @@ def search(
352
354
self ,
353
355
* ,
354
356
query : typing .Union ['QueryModel' , dict [str , typing .Any ]],
355
- require_all_on : typing . Union [ str , list [str ]] = None ,
357
+ require_all_on : str | list [str ] | None = None ,
356
358
) -> 'ESMCatalogModel' :
357
359
"""
358
360
Search for entries in the catalog.
@@ -398,13 +400,13 @@ def search(
398
400
class QueryModel (pydantic .BaseModel ):
399
401
"""A Pydantic model to represent a query to be executed against a catalog."""
400
402
401
- query : dict [pydantic .StrictStr , typing .Union [ typing . Any , list [typing .Any ] ]]
403
+ query : dict [pydantic .StrictStr , typing .Any | list [typing .Any ]]
402
404
columns : list [str ]
403
- require_all_on : typing . Optional [ typing . Union [ str , list [typing .Any ]]] = None
405
+ require_all_on : str | list [typing .Any ] | None = None
404
406
405
407
# TODO: Seem to be unable to modify fields in model_validator with
406
408
# validate_assignment=True since it leads to recursion
407
- model_config = ConfigDict (validate_default = True , validate_assignment = False )
409
+ model_config = ConfigDict (validate_assignment = False )
408
410
409
411
@pydantic .model_validator (mode = 'after' )
410
412
def validate_query (cls , model ):
@@ -424,7 +426,7 @@ def validate_query(cls, model):
424
426
raise ValueError (f'Column { key } not in columns { columns } ' )
425
427
_query = query .copy ()
426
428
for key , value in _query .items ():
427
- if isinstance (value , ( str , int , float , bool ) ) or value is None or value is pd .NA :
429
+ if isinstance (value , str | int | float | bool ) or value is None or value is pd .NA :
428
430
_query [key ] = [value ]
429
431
430
432
model .query = _query
0 commit comments