From d735d254a8b87914a21c1f3034c12374e02af324 Mon Sep 17 00:00:00 2001
From: XiaoXinYo <1104361313@qq.com>
Date: Mon, 29 May 2023 17:57:32 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96GenerateResponse?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 module/core.py | 49 +++++++++++++++++++++++++++----------------------
 1 file changed, 27 insertions(+), 22 deletions(-)

diff --git a/module/core.py b/module/core.py
index 0cea66b..12c47aa 100644
--- a/module/core.py
+++ b/module/core.py
@@ -19,29 +19,34 @@ class GenerateResponse:
     TYPE = Union[str, Response]
 
     def __init__(self):
-        self.response = {}
-        self.onlyJSON = False
-    
-    def json(self) -> TYPE:
-        responseJSON = json.dumps(self.response, ensure_ascii=False)
-        if self.stream:
+        self.code = 0
+        self.message = ''
+        self.data = None
+        self.streamFormat = False
+        self.streamResponse = False
+
+    def generate(self) -> TYPE:
+        responseJSON = json.dumps({
+            'code': self.code,
+            'message': self.message,
+            'data': self.data
+        }, ensure_ascii=False)
+        if self.streamFormat:
             return f'data: {responseJSON}\n\n'
+        if self.streamResponse:
+            return Response(f'data: {responseJSON}\n\n', media_type='text/event-stream')
         return Response(responseJSON, media_type='application/json')
 
-    def error(self, code: int, message: str, stream=False) -> TYPE:
-        self.response = {
-            'code': code,
-            'message': message,
-            'data': None
-        }
-        self.stream = stream
-        return self.json()
+    def error(self, code: int, message: str, streamFormat=False, streamResponse=False) -> TYPE:
+        self.code = code
+        self.message = message
+        self.streamFormat = streamFormat
+        self.streamResponse = streamResponse
+        return self.generate()
 
-    def success(self, data: Any, stream=False) -> TYPE:
-        self.response = {
-            'code': 200,
-            'message': 'success',
-            'data': data
-        }
-        self.stream = stream
-        return self.json()
\ No newline at end of file
+    def success(self, data: Any, streamFormat=False) -> TYPE:
+        self.code = 200
+        self.message = 'success'
+        self.data = data
+        self.streamFormat = streamFormat
+        return self.generate()
\ No newline at end of file