From f597262bddf94bd16a955bfca28ebc8b22523cd3 Mon Sep 17 00:00:00 2001
From: brb <brb@esoft.com>
Date: Thu, 13 Apr 2023 19:28:07 +0200
Subject: [PATCH] added definition of classes

---
 classes.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 classes.py

diff --git a/classes.py b/classes.py
new file mode 100644
index 0000000..fa355ac
--- /dev/null
+++ b/classes.py
@@ -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
\ No newline at end of file