@@ -51,9 +51,9 @@ def _build_columns( # pylint:disable=too-many-locals
51
51
check_name = getattr (field , "check_name" , None )
52
52
53
53
try :
54
- is_polars_dtype = inspect .isclass (
55
- annotation .raw_annotation
56
- ) and issubclass ( annotation . raw_annotation , pe . DataType )
54
+ is_polars_dtype = inspect .isclass (annotation . raw_annotation ) and issubclass (
55
+ annotation .raw_annotation , pe . DataType
56
+ )
57
57
except TypeError :
58
58
is_polars_dtype = False
59
59
@@ -87,9 +87,7 @@ def _build_columns( # pylint:disable=too-many-locals
87
87
or dtype
88
88
):
89
89
if check_name is False :
90
- raise SchemaInitError (
91
- f"'check_name' is not supported for { field_name } ."
92
- )
90
+ raise SchemaInitError (f"'check_name' is not supported for { field_name } ." )
93
91
94
92
column_kwargs = (
95
93
field .column_properties (
@@ -104,9 +102,7 @@ def _build_columns( # pylint:disable=too-many-locals
104
102
columns [field_name ] = Column (** column_kwargs )
105
103
106
104
else :
107
- raise SchemaInitError (
108
- f"Invalid annotation '{ field_name } : { annotation .raw_annotation } '."
109
- )
105
+ raise SchemaInitError (f"Invalid annotation '{ field_name } : { annotation .raw_annotation } '." )
110
106
111
107
return columns
112
108
@@ -149,9 +145,7 @@ def validate(
149
145
inplace : bool = False ,
150
146
) -> Union [LazyFrame [Self ], DataFrame [Self ]]:
151
147
"""%(validate_doc)s"""
152
- result = cls .to_schema ().validate (
153
- check_obj , head , tail , sample , random_state , lazy , inplace
154
- )
148
+ result = cls .to_schema ().validate (check_obj , head , tail , sample , random_state , lazy , inplace )
155
149
if isinstance (check_obj , pl .LazyFrame ):
156
150
return cast (LazyFrame [Self ], result )
157
151
else :
@@ -170,7 +164,7 @@ def to_json_schema(cls):
170
164
FastAPI integration.
171
165
"""
172
166
schema = cls .to_schema ()
173
-
167
+
174
168
# Define a mapping from Polars data types to JSON schema types
175
169
# This is more robust than string parsing
176
170
POLARS_TO_JSON_TYPE_MAP = {
@@ -183,52 +177,48 @@ def to_json_schema(cls):
183
177
pl .UInt16 : "integer" ,
184
178
pl .UInt32 : "integer" ,
185
179
pl .UInt64 : "integer" ,
186
-
187
180
# Float types
188
181
pl .Float32 : "number" ,
189
182
pl .Float64 : "number" ,
190
-
191
183
# Boolean type
192
184
pl .Boolean : "boolean" ,
193
-
194
185
# String types
195
186
pl .Utf8 : "string" ,
196
187
pl .String : "string" ,
197
-
198
188
# Date/Time types
199
189
pl .Date : "datetime" ,
200
190
pl .Datetime : "datetime" ,
201
191
pl .Time : "datetime" ,
202
192
pl .Duration : "datetime" ,
203
193
}
204
-
194
+
205
195
def map_dtype_to_json_type (dtype ):
206
196
"""
207
197
Map a Polars data type to a JSON schema type.
208
-
198
+
209
199
Args:
210
200
dtype: Polars data type
211
-
201
+
212
202
Returns:
213
203
str: JSON schema type string
214
204
"""
215
205
# First try the direct mapping
216
206
if dtype .__class__ in POLARS_TO_JSON_TYPE_MAP :
217
207
return POLARS_TO_JSON_TYPE_MAP [dtype .__class__ ]
218
-
208
+
219
209
# Fallback to string representation check for edge cases
220
210
dtype_str = str (dtype ).lower ()
221
- if ' float' in dtype_str :
211
+ if " float" in dtype_str :
222
212
return "number"
223
- elif ' int' in dtype_str :
213
+ elif " int" in dtype_str :
224
214
return "integer"
225
- elif ' bool' in dtype_str :
215
+ elif " bool" in dtype_str :
226
216
return "boolean"
227
- elif any (t in dtype_str for t in [' date' , ' time' , ' datetime' ]):
217
+ elif any (t in dtype_str for t in [" date" , " time" , " datetime" ]):
228
218
return "datetime"
229
219
else :
230
220
return "string"
231
-
221
+
232
222
properties = {}
233
223
for col_name , col_schema in schema .dtypes .items ():
234
224
json_type = map_dtype_to_json_type (col_schema .type )
0 commit comments