From 03905756ca6012d88d96a51b813a20fc50cb6552 Mon Sep 17 00:00:00 2001
From: Thomas Wade <gigatech55@gmail.com>
Date: Sat, 25 Apr 2020 03:06:13 +0930
Subject: [PATCH] Add search function

---
 FAAPI/FAAPI.py | 121 +++++++++++++++++++++++++++++++++++++++++++++++++
 setup.py       |   2 +-
 2 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/FAAPI/FAAPI.py b/FAAPI/FAAPI.py
index 9d6764f..b2c25a0 100644
--- a/FAAPI/FAAPI.py
+++ b/FAAPI/FAAPI.py
@@ -1,4 +1,5 @@
 from datetime import datetime
+from enum import Enum
 from typing import List
 
 import requests
@@ -444,6 +445,64 @@ class User:
         return self.username
 
 
+class OrderByEnum(Enum):
+    relevancy = 'relevancy'
+    date = 'date'
+    popularity = 'popularity'
+
+
+class OrderEnum(Enum):
+    ascending = 'asc'
+    descending = 'desc'
+
+
+class RangeEnum(Enum):
+    day = 'day'
+    three_days = '3days'
+    week = 'week'
+    month = 'month'
+    all_time = 'all'
+
+
+class MatchModeEnum(Enum):
+    all = 'all'
+    any = 'any'
+    extended = 'extended'
+
+
+class SearchOpts:
+    """
+    A container for FurAffinity search options.
+    """
+
+    def __init__(self, order_by: OrderByEnum = OrderByEnum.relevancy,
+                 order_direction: OrderEnum = OrderEnum.descending,
+                 range: RangeEnum = RangeEnum.all_time,
+                 match_mode: MatchModeEnum = MatchModeEnum.extended,
+                 rating_general: bool = True,
+                 rating_mature: bool = True,
+                 rating_adult: bool = True,
+                 type_art: bool = True,
+                 type_music: bool = True,
+                 type_flash: bool = True,
+                 type_story: bool = True,
+                 type_photo: bool = True,
+                 type_poetry: bool = True):
+        self.order_by: OrderByEnum = order_by
+        self.order_direction: OrderEnum = order_direction
+        self.range: RangeEnum = range
+        self.match_mode: MatchModeEnum = match_mode
+        self.rating_general: bool = rating_general
+        self.rating_mature: bool = rating_mature
+        self.rating_adult: bool = rating_adult
+        self.type_art: bool = type_art
+        self.type_music: bool = type_music
+        self.type_flash: bool = type_flash
+        self.type_story: bool = type_story
+        self.type_photo: bool = type_photo
+        self.type_poetry: bool = type_poetry
+
+
 class FAAPI:
     """
     A basic FurAffinity API instance.
@@ -567,6 +626,68 @@ class FAAPI:
             submission._api = self
         return submission
 
+    def search(self, search: str, results: int = 48, opts: SearchOpts = SearchOpts()) -> List[Submission]:
+        """
+        Searches FurAffinity with the given string.
+        Results are
+        :param search: Search term
+        :param results: Number of results to return
+        :param opts: Search options
+        :return: List containing results
+        """
+        url = FA_BASE_URL + '/search/'
+        form_data = {
+            'page': 1,
+            'q': search,
+            'do_search': 'Search',
+            'order-by': opts.order_by,
+            'order-direction': opts.order_direction,
+            'range': opts.range,
+            'mode': opts.match_mode
+        }
+
+        # Set optional keys since they're judged on presence, not value
+        if opts.rating_general: form_data['rating-general'] = 'on'
+        if opts.rating_mature: form_data['rating-mature'] = 'on'
+        if opts.rating_adult: form_data['rating-adult'] = 'on'
+        if opts.type_art: form_data['type-art'] = 'on'
+        if opts.type_music: form_data['type-music'] = 'on'
+        if opts.type_flash: form_data['type-flash'] = 'on'
+        if opts.type_story: form_data['type-story'] = 'on'
+        if opts.type_photo: form_data['type-photo'] = 'on'
+        if opts.type_poetry: form_data['type-poetry'] = 'on'
+
+        submissions = []
+        while len(submissions) < results:
+            # Get our search results
+            res = self._req.post(url, data=form_data)
+            if res.status_code >= 400:
+                raise Exception('Got status code {} when trying to get URL {}'.format(res.status_code, url))
+            soup = BeautifulSoup(res.content, 'lxml')
+
+            # Get all the submissions on this page
+            items = soup.select('.gallery figure')
+            for item in items:
+                sub = self.get_submission(int(item.get('id')[4:]))  # IDs look like 'sid-35908275', so we just skip the 'sid-'
+
+                # Try cache the preview URL early
+                try:
+                    sub._preview_url = 'https:' + item.select('img')[0].get('src')
+                    sub._has_preview = True
+                except IndexError:
+                    sub._has_preview = False
+
+                submissions.append(sub)
+
+            if not soup.select('input[name=next_page]:not(.hidden)'):
+                # No further pages, stop searching
+                break
+
+            # Increment our page number
+            form_data['page'] = form_data['page'] + 1
+
+        return submissions[:results]
+
     def __getstate__(self):
         # Get a copy of the object state without the requests session
         state = self.__dict__.copy()
diff --git a/setup.py b/setup.py
index ff99c5b..dcfed98 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='FAAPI',
-    version='0.5.4',
+    version='0.6.0',
     packages=find_packages(),
     url='https://tem.party/gitea/tom/FAAPI',
     license='WTFPL',