added definition of tweet class
This commit is contained in:
@@ -0,0 +1,40 @@
|
||||
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]]]
|
||||
video: Optional[str]
|
||||
|
||||
class Config:
|
||||
arbitrary_types_allowed = True
|
||||
|
||||
@validator('images', each_item=True)
|
||||
def parse_image(cls, v):
|
||||
if isinstance(v, BytesIO):
|
||||
v = Image.open(v)
|
||||
return v
|
||||
|
||||
def parse_images(self):
|
||||
for i in range(len(self.images)):
|
||||
self.images[i] = Image.open(self.images[i])
|
||||
|
||||
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