added definition of classes
This commit is contained in:
+41
@@ -0,0 +1,41 @@
|
|||||||
|
from pydantic import BaseModel, validator
|
||||||
|
from typing import Union, List, Optional
|
||||||
|
from datetime import datetime
|
||||||
|
from PIL import Image
|
||||||
|
from io import BytesIO
|
||||||
|
|
||||||
|
class Tweet(BaseModel):
|
||||||
|
account: str
|
||||||
|
url: str
|
||||||
|
time: datetime
|
||||||
|
text: str
|
||||||
|
images: Optional[Union[List[Image.Image], List[BytesIO], List[bytes]]]
|
||||||
|
video: Optional[str]
|
||||||
|
|
||||||
|
@validator('images')
|
||||||
|
def parse_images(cls, v):
|
||||||
|
parsed_elem_list = list()
|
||||||
|
for elem in v:
|
||||||
|
if isinstance(elem, bytes):
|
||||||
|
elem = BytesIO(elem)
|
||||||
|
if isinstance(elem, BytesIO):
|
||||||
|
elem = Image.open(elem)
|
||||||
|
parsed_elem_list.append(elem)
|
||||||
|
return parsed_elem_list
|
||||||
|
|
||||||
|
class Config: # allows subtypes to be e.g. PIL.Image.Image without error
|
||||||
|
arbitrary_types_allowed = True
|
||||||
|
|
||||||
|
def as_db_dict(self):
|
||||||
|
"""Convert to database dictionary"""
|
||||||
|
d = self.dict()
|
||||||
|
d['time'] = d['time'].strftime('%Y-%m-%d %H:%M:%S')
|
||||||
|
img_list = list()
|
||||||
|
for img in d['images']:
|
||||||
|
if isinstance(img, BytesIO):
|
||||||
|
img = Image.open(img)
|
||||||
|
buf = BytesIO()
|
||||||
|
img.save(buf, format='PNG')
|
||||||
|
img_list.append(buf.getvalue())
|
||||||
|
d['images'] = img_list
|
||||||
|
return d
|
||||||
Reference in New Issue
Block a user