""" 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.dataclasses 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: """Check 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: """Generate 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): # noqa: D101 pass class MissingRequiredParameterError(SzurubooruException): # noqa: D101 pass class InvalidParameterError(SzurubooruException): # noqa: D101 pass class IntegrityError(SzurubooruException): # noqa: D101 pass class SearchError(SzurubooruException): # noqa: D101 pass class AuthError(SzurubooruException): # noqa: D101 pass class PostNotFoundError(SzurubooruException): # noqa: D101 pass class PostAlreadyFeaturedError(SzurubooruException): # noqa: D101 pass class PostAlreadyUploadedError(SzurubooruException): # noqa: D101 pass class InvalidPostIdError(SzurubooruException): # noqa: D101 pass class InvalidPostSafetyError(SzurubooruException): # noqa: D101 pass class InvalidPostSourceError(SzurubooruException): # noqa: D101 pass class InvalidPostContentError(SzurubooruException): # noqa: D101 pass class InvalidPostRelationError(SzurubooruException): # noqa: D101 pass class InvalidPostNoteError(SzurubooruException): # noqa: D101 pass class InvalidPostFlagError(SzurubooruException): # noqa: D101 pass class InvalidFavoriteTargetError(SzurubooruException): # noqa: D101 pass class InvalidCommentIdError(SzurubooruException): # noqa: D101 pass class CommentNotFoundError(SzurubooruException): # noqa: D101 pass class EmptyCommentTextError(SzurubooruException): # noqa: D101 pass class InvalidScoreTargetError(SzurubooruException): # noqa: D101 pass class InvalidScoreValueError(SzurubooruException): # noqa: D101 pass class TagCategoryNotFoundError(SzurubooruException): # noqa: D101 pass class TagCategoryAlreadyExistsError(SzurubooruException): # noqa: D101 pass class TagCategoryIsInUseError(SzurubooruException): # noqa: D101 pass class InvalidTagCategoryNameError(SzurubooruException): # noqa: D101 pass class InvalidTagCategoryColorError(SzurubooruException): # noqa: D101 pass class TagNotFoundError(SzurubooruException): # noqa: D101 pass class TagAlreadyExistsError(SzurubooruException): # noqa: D101 pass class TagIsInUseError(SzurubooruException): # noqa: D101 pass class InvalidTagNameError(SzurubooruException): # noqa: D101 pass class InvalidTagRelationError(SzurubooruException): # noqa: D101 pass class InvalidTagCategoryError(SzurubooruException): # noqa: D101 pass class InvalidTagDescriptionError(SzurubooruException): # noqa: D101 pass class UserNotFoundError(SzurubooruException): # noqa: D101 pass class UserAlreadyExistsError(SzurubooruException): # noqa: D101 pass class InvalidUserNameError(SzurubooruException): # noqa: D101 pass class InvalidEmailError(SzurubooruException): # noqa: D101 pass class InvalidPasswordError(SzurubooruException): # noqa: D101 pass class InvalidRankError(SzurubooruException): # noqa: D101 pass class InvalidAvatarError(SzurubooruException): # noqa: D101 pass class ProcessingError(SzurubooruException): # noqa: D101 pass class ValidationError(SzurubooruException): # noqa: D101 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: """Map 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: """Map 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=list(map(lambda a: a["names"][0], c["suggestions"])), ) def tag_exists(self, name: str) -> bool: """Check 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: """Check 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]: """List 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]: """List 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: """Delete 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: """Remove 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]]: """List 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))