From ae3a32589a39e59fc05aad65b6d5ed0166ccfa39 Mon Sep 17 00:00:00 2001
From: HibiKier <45528451+HibiKier@users.noreply.github.com>
Date: Mon, 19 Jul 2021 15:23:31 +0800
Subject: [PATCH] Update data_source.py

---
 plugins/send_setu/data_source.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/plugins/send_setu/data_source.py b/plugins/send_setu/data_source.py
index 206aae3d..7e4fa4be 100644
--- a/plugins/send_setu/data_source.py
+++ b/plugins/send_setu/data_source.py
@@ -117,10 +117,15 @@ async def search_online_setu(url: str):
         return '\n图片被小怪兽恰掉啦..!QAQ', -1
 
 
-def get_setu(index: str, setu_data: dict, tag: str = None):
-    if not os.path.exists(IMAGE_PATH + path):
-        os.mkdir(IMAGE_PATH + path)
-    length = len(os.listdir(IMAGE_PATH + path))
+def get_setu(index: str, setu_data: dict, num: int = 1, tag: str = None, r18: bool = False):
+    if r18:
+        if not os.path.exists(IMAGE_PATH + '_r18'):
+            os.mkdir(IMAGE_PATH + '_r18')
+        length = len(os.listdir(IMAGE_PATH + '_r18'))
+    else:
+        if not os.path.exists(IMAGE_PATH + path):
+            os.mkdir(IMAGE_PATH + path)
+        length = len(os.listdir(IMAGE_PATH + path))
     if length == 0:
         return None, None
     if not index: