Utilizes model_validate instead of parse_obj

Deprecation warnings -- this is the new/better way to do it. Also makes pydantic>1 for the model_validate methods
DAGWorks-Inc · Dec 11, 2024 · 4c7ccbd · 4c7ccbd
1 parent 1563a46
commit 4c7ccbd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/burr/tracking/server/schema.py b/burr/tracking/server/schema.py
@@ -93,30 +93,30 @@ def from_logs(log_lines: List[bytes]) -> List["Step"]:
             json_line = safe_json_load(line)
             # TODO -- make these into constants
             if json_line["type"] == "begin_entry":
-                begin_step = BeginEntryModel.parse_obj(json_line)
+                begin_step = BeginEntryModel.model_validate(json_line)
                 steps_by_sequence_id[begin_step.sequence_id].step_start_log = begin_step
             elif json_line["type"] == "end_entry":
-                step_end_log = EndEntryModel.parse_obj(json_line)
+                step_end_log = EndEntryModel.model_validate(json_line)
                 steps_by_sequence_id[step_end_log.sequence_id].step_end_log = step_end_log
             elif json_line["type"] == "begin_span":
-                span = BeginSpanModel.parse_obj(json_line)
+                span = BeginSpanModel.model_validate(json_line)
                 spans_by_id[span.span_id] = PartialSpan(
                     begin_entry=span,
                     end_entry=None,
                 )
             elif json_line["type"] == "end_span":
-                end_span = EndSpanModel.parse_obj(json_line)
+                end_span = EndSpanModel.model_validate(json_line)
                 span = spans_by_id[end_span.span_id]
                 span.end_entry = end_span
             elif json_line["type"] == "attribute":
-                attribute = AttributeModel.parse_obj(json_line)
+                attribute = AttributeModel.model_validate(json_line)
                 attributes_by_step[attribute.action_sequence_id].append(attribute)
             elif json_line["type"] in ["begin_stream", "first_item_stream", "end_stream"]:
                 streaming_event = {
                     "begin_stream": InitializeStreamModel,
                     "first_item_stream": FirstItemStreamModel,
                     "end_stream": EndStreamModel,
-                }[json_line["type"]].parse_obj(json_line)
+                }[json_line["type"]].model_validate(json_line)
                 steps_by_sequence_id[streaming_event.sequence_id].streaming_events.append(
                     streaming_event
                 )

diff --git a/pyproject.toml b/pyproject.toml
@@ -79,7 +79,7 @@ documentation = [
 ]
 
 tracking-client = [
-  "pydantic"
+  "pydantic>1"
 ]
 
 tracking-client-s3 = [