""" Szurubooru Library ~~~~~~~~~~~~~~~~~~ Szurubooru Library is used for communicating with a Szurubooru instance through python. :copyright (c) 2023 Colin Wilk. :license: MIT, see LICENSE for more details. """ from lib.tag import Tag import requests as r from typing import List import json from urllib.parse import quote import math import functools class SzurubooruException(Exception): """General Error returned from the Szurubooru REST API. The Szurubooru API returns errors together with an http code in a form like this: .. code-block:: JSON { "name": "Name of the error, e.g. 'PostNotFoundError'", "title": "Generic title of error message, e.g. 'Not found'", "description": "Detailed description of what went wrong, e.g. 'User `rr-` not found." } We map this json response to the SzurubooruException class or more specific children such as MissingRequiredFileError with the fields / args that are provided from the json response. Args: name: Name of the exception as defined by the Szurubooru REST API The name is also used in more specific Exceptions such as MissingRequiredFileError(SzurubooruException). title: Generic, human readable title of the error message. description: More detailed description of what went wrong. Often includes input from the user that caused the error. """ name: str title: str description: str def __init__(self, name: str, title: str, description: str) -> None: self.name = name self.title = title self.description = description super().__init__(description) @staticmethod def response_is_exception(res: r.Response) -> bool: """Checks if the Szurubooru REST API response is an error or not. Args: res (r.Response): the Response of the request from the requests module. Returns: bool: True if response is an error and False if it is not. """ return res.status_code != 200 @staticmethod def map_exception(res: r.Response) -> Exception: """Generates an exception for an error response from the Szurubooru REST API. Args: res (r.Response): the Response from the Szurubooru REST API that should be mapped to a Python Exception. Raises: LookupError: Thrown when the Szurubooru REST API sent an error type (name) that could not be mapped to one of the SzurubooruException Python Exceptions. This can happen when there was an unexpected e.g. an Internal Server Error in the Szurubooru REST API. Returns: Exception: Returns a more specific exception. """ content: dict[str, str] = json.loads(res.content) name: str = content["name"] title: str = content["title"] description: str = content["description"] match name: case "MissingRequiredFileError": return MissingRequiredFileError(name, title, description) case "MissingRequiredParameterError": return MissingRequiredParameterError(name, title, description) case "InvalidParameterError": return InvalidParameterError(name, title, description) case "IntegrityError": return IntegrityError(name, title, description) case "SearchError": return SearchError(name, title, description) case "AuthError": return AuthError(name, title, description) case "PostNotFoundError": return PostNotFoundError(name, title, description) case "PostAlreadyFeaturedError": return PostAlreadyFeaturedError(name, title, description) case "PostAlreadyUploadedError": return PostAlreadyUploadedError(name, title, description) case "InvalidPostIdError": return InvalidPostIdError(name, title, description) case "InvalidPostSafetyError": return InvalidPostSafetyError(name, title, description) case "InvalidPostSourceError": return InvalidPostSourceError(name, title, description) case "InvalidPostContentError": return InvalidPostContentError(name, title, description) case "InvalidPostRelationError": return InvalidPostRelationError(name, title, description) case "InvalidPostNoteError": return InvalidPostNoteError(name, title, description) case "InvalidPostFlagError": return InvalidPostFlagError(name, title, description) case "InvalidFavoriteTargetError": return InvalidFavoriteTargetError(name, title, description) case "InvalidCommentIdError": return InvalidCommentIdError(name, title, description) case "CommentNotFoundError": return CommentNotFoundError(name, title, description) case "EmptyCommentTextError": return EmptyCommentTextError(name, title, description) case "InvalidScoreTargetError": return InvalidScoreTargetError(name, title, description) case "InvalidScoreValueError": return InvalidScoreValueError(name, title, description) case "TagCategoryNotFoundError": return TagCategoryNotFoundError(name, title, description) case "TagCategoryAlreadyExistsError": return TagCategoryAlreadyExistsError(name, title, description) case "TagCategoryIsInUseError": return TagCategoryIsInUseError(name, title, description) case "InvalidTagCategoryNameError": return InvalidTagCategoryNameError(name, title, description) case "InvalidTagCategoryColorError": return InvalidTagCategoryColorError(name, title, description) case "TagNotFoundError": return TagNotFoundError(name, title, description) case "TagAlreadyExistsError": return TagAlreadyExistsError(name, title, description) case "TagIsInUseError": return TagIsInUseError(name, title, description) case "InvalidTagNameError": return InvalidTagNameError(name, title, description) case "InvalidTagRelationError": return InvalidTagRelationError(name, title, description) case "InvalidTagCategoryError": return InvalidTagCategoryError(name, title, description) case "InvalidTagDescriptionError": return InvalidTagDescriptionError(name, title, description) case "UserNotFoundError": return UserNotFoundError(name, title, description) case "UserAlreadyExistsError": return UserAlreadyExistsError(name, title, description) case "InvalidUserNameError": return InvalidUserNameError(name, title, description) case "InvalidEmailError": return InvalidEmailError(name, title, description) case "InvalidPasswordError": return InvalidPasswordError(name, title, description) case "InvalidRankError": return InvalidRankError(name, title, description) case "InvalidAvatarError": return InvalidAvatarError(name, title, description) case "ProcessingError": return ProcessingError(name, title, description) case "ValidationError": return ValidationError(name, title, description) case _: raise LookupError(f'Unknown SzurubooruException: "{name}"') class MissingRequiredFileError(SzurubooruException): pass class MissingRequiredParameterError(SzurubooruException): pass class InvalidParameterError(SzurubooruException): pass class IntegrityError(SzurubooruException): pass class SearchError(SzurubooruException): pass class AuthError(SzurubooruException): pass class PostNotFoundError(SzurubooruException): pass class PostAlreadyFeaturedError(SzurubooruException): pass class PostAlreadyUploadedError(SzurubooruException): pass class InvalidPostIdError(SzurubooruException): pass class InvalidPostSafetyError(SzurubooruException): pass class InvalidPostSourceError(SzurubooruException): pass class InvalidPostContentError(SzurubooruException): pass class InvalidPostRelationError(SzurubooruException): pass class InvalidPostNoteError(SzurubooruException): pass class InvalidPostFlagError(SzurubooruException): pass class InvalidFavoriteTargetError(SzurubooruException): pass class InvalidCommentIdError(SzurubooruException): pass class CommentNotFoundError(SzurubooruException): pass class EmptyCommentTextError(SzurubooruException): pass class InvalidScoreTargetError(SzurubooruException): pass class InvalidScoreValueError(SzurubooruException): pass class TagCategoryNotFoundError(SzurubooruException): pass class TagCategoryAlreadyExistsError(SzurubooruException): pass class TagCategoryIsInUseError(SzurubooruException): pass class InvalidTagCategoryNameError(SzurubooruException): pass class InvalidTagCategoryColorError(SzurubooruException): pass class TagNotFoundError(SzurubooruException): pass class TagAlreadyExistsError(SzurubooruException): pass class TagIsInUseError(SzurubooruException): pass class InvalidTagNameError(SzurubooruException): pass class InvalidTagRelationError(SzurubooruException): pass class InvalidTagCategoryError(SzurubooruException): pass class InvalidTagDescriptionError(SzurubooruException): pass class UserNotFoundError(SzurubooruException): pass class UserAlreadyExistsError(SzurubooruException): pass class InvalidUserNameError(SzurubooruException): pass class InvalidEmailError(SzurubooruException): pass class InvalidPasswordError(SzurubooruException): pass class InvalidRankError(SzurubooruException): pass class InvalidAvatarError(SzurubooruException): pass class ProcessingError(SzurubooruException): pass class ValidationError(SzurubooruException): pass class Szurubooru: """Client class for communicating with the Szurubooru instance. Raises: SzurubooruException: When the Szurubooru REST API returns any kinds of errors. See: SzurubooruException for mor information. """ url: str token: str verify: bool header: dict[str, str] # Session Stats requests_get: int = 0 requests_post: int = 0 requests_put: int = 0 requests_delete: int = 0 def __init__(self, url: str, token: str, verify: bool = True): self.url = url self.token = token self.verify = verify self.header = { "Accept": "application/json", "Authorization": f"Token {self.token}", } @staticmethod def __raise_szurubooru_exception(func): @functools.wraps(func) def wrapper(self, **kwargs): res = func(self, **kwargs) if SzurubooruException.response_is_exception(res): raise SzurubooruException.map_exception(res) return res return wrapper @staticmethod def __count_stats(get: int = 0, post: int = 0, put: int = 0, delete: int = 0): def function_wrapper(func): @functools.wraps(func) def wrapper(self, **kwargs): self.requests_get += get self.requests_post += post self.requests_put += put self.requests_delete += delete return func(self, **kwargs) return wrapper return function_wrapper @__count_stats(get=1) @__raise_szurubooru_exception def __get(self, path: str) -> r.Response: return r.get( f"{self.url}{path}", timeout=15, verify=self.verify, headers=self.header ) @__count_stats(post=1) @__raise_szurubooru_exception def __post(self, path: str, params: dict) -> r.Response: return r.post( f"{self.url}{path}", timeout=15, verify=self.verify, headers=self.header, data=json.dumps(params), ) @__count_stats(put=1) @__raise_szurubooru_exception def __put(self, path: str, params: dict) -> r.Response: return r.put( f"{self.url}{path}", timeout=15, verify=self.verify, headers=self.header, data=json.dumps(params), ) @__count_stats(delete=1) @__raise_szurubooru_exception def __delete(self, path: str, params: dict) -> r.Response: return r.delete( f"{self.url}{path}", timeout=15, verify=self.verify, headers=self.header, data=json.dumps(params), ) def __map_tag_resource_response(self, res: r.Response, name: str = "") -> Tag: """Maps the requests response from the Szurubooru REST API to the Tag dataclass. This function is similar to __map_tag_resource and calls it internally. Instead of receiving a dictionary if the Szurubooru Tag it operates on the requests Response of the same form. To Illustrate, both code yield the same result: >>> __map_tag_resource_response(res, name) >>> __map_tag_resource(json.loads(res.content), name) Args: res (r.Response): The requests Response from an Szurubooru REST API endpoint that returns a Szurubooru tag resource in a form like: .. code-block:: JSON { "version": , "names": , "category": , "implications": , "suggestions": , "creationTime": , "lastEditTime": , "usages": , "description": } name (str, optional): If you provide a name here, the mapper will use that name for the Tag dataclass instead of fetching it from the alias names. This is useful if you have multiple aliases or names in Szurubooru. Returns: Tag: A Tag dataclass. The version parameter is included if the Szurubooru REST API response contained a version field, otherwise it defaults to -1. The name will be the first name returned by the Szurubooru REST API or the name that was passed to the function as an optional parameter. """ return self.__map_tag_resource(json.loads(res.content), name) def __map_tag_resource(self, c: dict, name: str = "") -> Tag: """Maps the dict from the Szurubooru REST API to the Tag dataclass. This function is similar to __map_tag_resource_response and is called by it. Instead of receiving a requests Response it operates on the dictionary of the same form. To Illustrate, both code yield the same result: >>> __map_tag_resource_response(res, name) >>> __map_tag_resource(json.loads(res.content), name) Args: c (dict): The dict object coming from an Szurubooru REST API Szurubooru tag resource in a form like: .. code-block:: JSON { "version": , "names": , "category": , "implications": , "suggestions": , "creationTime": , "lastEditTime": , "usages": , "description": } name (str, optional): If you provide a name here, the mapper will use that name for the Tag dataclass instead of fetching it from the alias names. This is useful if you have multiple aliases or names in Szurubooru. Returns: Tag: A Tag dataclass. The version parameter is included if the Szurubooru REST API response contained a version field, otherwise it defaults to -1. The name will be the first name returned by the Szurubooru REST API or the name that was passed to the function as an optional parameter. """ return Tag( name=name if name != "" else c["names"][0], version=c.get("version", -1), category=c["category"], description=c["description"], implications=list(map(lambda a: a["names"][0], c["implications"])), suggestions=c["suggestions"], ) def tag_exists(self, name: str) -> bool: """Checks if the tag exists on Szurubooru. Args: name (str): Name of the tag. Returns: bool: True if tag already exists. False if there was no matching tag on Szurubooru. """ try: res = self.__get(path=f"/api/tag/{quote(name)}") except TagNotFoundError: return False except Exception as e: raise e return res.status_code == 200 def tag_len(self) -> int: """Checks how many tags there are on the Szurubooru instance. Returns: int: The number of tags that exist on the Szurubooru instance """ return json.loads(self.__get(path="/api/tags").content)["total"] def tag_list(self, limit: int = 100, offset: int = 0, query: str = "") -> List[Tag]: """Lists tags on the Szurubooru instance. Args: limit (int, optional): Limits the amount of tags displayed on one page. Setting a smaller limit may help with query speed. Usually the maximum allowed limit is 100. Defaults to 100. offset (int, optional): Offset where to start the listing. Required for Paging the data when searching through many tags. If you have a page size of 100 you need to set an offset of 100 to see results on the 2. page. So your offset should be PAGE_NUM * LIMIT. Defaults to 0. query (str, optional): You can use a custom query to filter through tags better. If you use an empty string (the default) you will search through all tags on the Szurubooru instance. Defaults to ''. Returns: List[Tag]: List of Tag objects in the current page. """ path = f"/api/tags?offset={offset}&limit={limit}&query={quote(query)}" results = json.loads(self.__get(path=path).content)["results"] return list(map(self.__map_tag_resource, results)) def tag_list_all(self, query: str = "") -> List[Tag]: """Lists every available Tag (matching the query) on the Szurubooru instance. WARNING: This method calls the Szurubooru REST API multiple times depending on how many tags there are. It can take a while and can put significant load on the Szurubooru instance. Use with caution. In essence this method does the paging for you in tag_list and returns the full batch of Tags. Args: query (str, optional): You can use a custom query to filter through tags better. If you use an empty string (the default) you will search through all tags on the Szurubooru instance. Defaults to ''. Returns: List[Tag]: List of all Tag objects matching the query on the Szurubooru instance. """ tags = [] for i in range(0, math.ceil(self.tag_len() / 100)): tags.extend(self.tag_list(offset=100 * i, limit=100, query=query)) return tags def tag_get(self, name: str) -> Tag: """Get a tag form the Szurubooru instance by name. Args: name (str): Name of the tag. Raises: TagNotFoundError: Tag could not be found on the Szurubooru instance. Returns: Tag: Tag object of the requested tag. """ res = self.__get(path=f"/api/tag/{quote(name)}") return self.__map_tag_resource_response(res, name=name) def tag_create(self, tag: Tag) -> Tag: """Create a tag on the Szurubooru instance. Args: tag (Tag): Tag dataclass representing the Tag that should be created. Returns: Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API representing the created tag as it was created on the Szurubooru instance. """ res = self.__post( path="/api/tags", params={ "names": [tag.name], "category": tag.category, "description": tag.description, "implications": tag.implications, "suggestions": tag.suggestions, }, ) return self.__map_tag_resource_response(res, name=tag.name) def tag_update(self, tag: Tag) -> Tag: """Update a tag on the Szurubooru instance. Args: tag (Tag): Tag dataclass representing the Tag that should be updated. The name is the identifying parameter for the Szurubooru instance. Returns: Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API representing the updated version of the tag as it is now on the Szurubooru instance. """ res = self.__put( path=f"/api/tag/{quote(tag.name)}", params={ "version": tag.version, "category": tag.category, "description": tag.description, "implications": tag.implications, "suggestions": tag.suggestions, }, ) return self.__map_tag_resource_response(res, name=tag.name) def tag_delete(self, tag: Tag) -> None: """Deletes the given tag on the Szurubooru instance. Args: tag (Tag): Tag dataclass of the tag that should be reledte. Needs to contain the current version number. """ self.__delete( path=f"/api/tag/{quote(tag.name)}", params={"version": tag.version} ) return def tag_merge(self, tag_to_merge: Tag, tag_to_merge_into: Tag) -> Tag: """Removes tag_to_merge and merges all uses into tag_to_merge_into. Removes source tag and merges all of its usages, suggestions and implications to the target tag. Other tag properties such as category and aliases do not get transferred and are discarded. Args: tag_to_merge (Tag): Tag dataclass representing the Tag that will be removed and merged into the second Tag. The version parameter is required. tag_to_merge_into (Tag): Tag dataclass representing the Tag that the first tag will be merged into. The version parameter is required. Returns: Tag: Resulting Tag dataclass that was returned by the Szurubooru REST API representing the new (merged) tag on the Szurubooru instance. """ res = self.__post( path="/api/tag-merge", params={ "removeVersion": tag_to_merge.version, "remove": tag_to_merge.name, "mergeToVersion": tag_to_merge_into.version, "mergeTo": tag_to_merge_into.name, }, ) return self.__map_tag_resource_response(res, name=tag_to_merge_into.name) def tag_list_siblings(self, tag: Tag) -> List[dict[str, int]]: """Lists siblings of the given Tag Lists siblings of given tag, e.g. tags that were used in the same posts as the given tag. `occurrences` field signifies how many times a given sibling appears with given tag. Results are sorted by occurrences count and the list is truncated to the first 50 elements. Doesn't use paging. Args: tag (Tag): Tag that will queried. Only uses the tag.name parameter. Returns: List[dict[str, int]]: Returns a List of dictionaries with the primary tag name as the key and the number of occurrences as the value. """ res = self.__get(path=f"/api/tag-siblings/{tag.name}") results = json.loads(res.content)["results"] return list(map(lambda a: {a["tag"]["names"][0]: a["occurrences"]}, results))