tt-flake/pkgs/smi/log.patch

53 lines
1.7 KiB
Diff
Raw Permalink Normal View History

2024-01-31 05:20:40 +00:00
diff --git a/tt_smi/log.py b/tt_smi/log.py
index e7ae8fd..083bdb7 100644
--- a/tt_smi/log.py
+++ b/tt_smi/log.py
@@ -10,8 +10,8 @@ import base64
import inspect
import datetime
from pathlib import Path
-from typing import Any, Union, List, TypeVar, Generic
-from pydantic import BaseModel
+from typing import Any, Union, List, TypeVar, Generic, Optional
+from pydantic import BaseModel, create_model
from pydantic.fields import Field
@@ -38,26 +38,24 @@ class Date(datetime.datetime):
def optional(*fields):
- """Decorator function used to modify a pydantic model's fields to all be optional.
- Alternatively, you can also pass the field names that should be made optional as arguments
- to the decorator.
- Taken from https://github.com/samuelcolvin/pydantic/issues/1223#issuecomment-775363074
- """
-
- def dec(_cls):
+ def dec(cls):
+ fields_dict = {}
for field in fields:
- _cls.__fields__[field].required = False
- _cls.__fields__[field].default = None
- return _cls
+ field_info = cls.__annotations__.get(field)
+ if field_info is not None:
+ fields_dict[field] = (Optional[field_info], None)
+ OptionalModel = create_model(cls.__name__, **fields_dict)
+ OptionalModel.__module__ = cls.__module__
+
+ return OptionalModel
if fields and inspect.isclass(fields[0]) and issubclass(fields[0], BaseModel):
cls = fields[0]
- fields = cls.__fields__
+ fields = cls.__annotations__
return dec(cls)
return dec
-
def type_to_mapping(type: Any):
if issubclass(type, float):
return {"type": "float"}