python-ballchasing 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

ballchasing/__init__.py

@@ -65,3 +65,10 @@ from .constants import (
     TeamIdentification,
     Map
 )
+
+from .typed import (
+    ShallowReplay,
+    ShallowGroup,
+    DeepReplay,
+    DeepGroup
+)

ballchasing/api.py

@@ -1,14 +1,17 @@
 import os
 import time
 from datetime import datetime
-from typing import Optional, Iterator, Union, List
+from pathlib import Path
+from typing import Optional, Union, List, BinaryIO, Iterator
 from urllib.parse import parse_qs, urlparse
 
-from requests import sessions, Response, ConnectionError
+from requests import sessions, Response, ConnectionError, HTTPError
 
 from ballchasing.constants import GroupSortBy, SortDir, AnyPlaylist, AnyMap, AnySeason, AnyRank, AnyReplaySortBy, \
     AnySortDir, AnyVisibility, AnyGroupSortBy, AnyPlayerIdentification, AnyTeamIdentification, AnyMatchResult
-from .util import rfc3339, parse_replay_stats
+from ballchasing.typed import DeepReplay, ShallowReplay, DeepGroup, ShallowGroup
+from .typed.shared import BaseGroup, BasicGroup
+from .util import to_rfc3339, parse_replay_stats
 
 DEFAULT_URL = "https://ballchasing.com/api"
 
@@ -18,12 +21,16 @@ class BallchasingApi:
     Class for communication with ballchasing.com API (https://ballchasing.com/doc/api)
     """
 
-    def __init__(self,
-                 auth_key: str,
-                 sleep_time_on_rate_limit: Optional[float] = None,
-                 print_on_rate_limit: bool = False,
-                 base_url=None,
-                 do_initial_ping=True):
+    def __init__(
+            self,
+            auth_key: str,
+            *,
+            sleep_time_on_rate_limit: Optional[float] = None,
+            print_on_rate_limit: bool = False,
+            base_url=None,
+            do_initial_ping=True,
+            typed=False,
+    ):
         """
 
         :param auth_key: authentication key for API calls.
@@ -49,6 +56,7 @@ class BallchasingApi:
         else:
             self.sleep_time_on_rate_limit = sleep_time_on_rate_limit
         self.print_on_rate_limit = print_on_rate_limit
+        self.typed = typed
 
     @property
     def steam_name(self):
@@ -74,11 +82,12 @@ class BallchasingApi:
             self.ping()
         return self._ping_result.get("quota")
 
-    def _request(self,
-                 url_or_endpoint: str,
-                 method: str,
-                 **params
-                 ) -> Response:
+    def _request(
+            self,
+            url_or_endpoint: str,
+            method: str,
+            **params
+    ) -> Response:
         """
         Helper method for all requests.
 
@@ -92,7 +101,8 @@ class BallchasingApi:
         headers = {"Authorization": self.auth_key}
         url = f"{self.base_url}{url_or_endpoint}" if url_or_endpoint.startswith("/") else url_or_endpoint
         max_retries = 8
-        for retries in range(max_retries):
+        retries = 0
+        while True:
             try:
                 r: Response = self._session.request(method=method, url=url, headers=headers, **params)
                 if 200 <= r.status_code < 300:
@@ -108,13 +118,14 @@ class BallchasingApi:
                     elif self.sleep_time_on_rate_limit:
                         time.sleep(self.sleep_time_on_rate_limit)
                 else:
-                    r.raise_for_status()  # Raise an error for any other status code
+                    r.raise_for_status()  # Raise an error for any other status code'
             except ConnectionError as e:
                 if retries >= max_retries - 1:
                     raise e
                 s = 2 ** retries
                 print(f"Connection error, trying again in {s} seconds...")
                 time.sleep(s)
+                retries += 1
 
     def ping(self) -> dict:
         """
@@ -130,28 +141,56 @@ class BallchasingApi:
         self._ping_result = result
         return result
 
-    def get_replays(self,
-                    title: Optional[str] = None,
-                    player_name: Optional[Union[str, List[str]]] = None,
-                    player_id: Optional[Union[str, List[str]]] = None,
-                    playlist: Optional[Union[AnyPlaylist, List[AnyPlaylist]]] = None,
-                    season: Optional[Union[AnySeason, List[AnySeason]]] = None,
-                    match_result: Optional[Union[AnyMatchResult, List[AnyMatchResult]]] = None,
-                    min_rank: Optional[AnyRank] = None,
-                    max_rank: Optional[AnyRank] = None,
-                    pro: Optional[bool] = None,
-                    uploader: Optional[str] = None,
-                    group_id: Optional[Union[str, List[str]]] = None,
-                    map_id: Optional[Union[AnyMap, List[AnyMap]]] = None,
-                    created_before: Optional[Union[str, datetime]] = None,
-                    created_after: Optional[Union[str, datetime]] = None,
-                    replay_after: Optional[Union[str, datetime]] = None,
-                    replay_before: Optional[Union[str, datetime]] = None,
-                    count: int = 150,
-                    sort_by: Optional[AnyReplaySortBy] = None,
-                    sort_dir: AnySortDir = SortDir.DESCENDING,
-                    deep: bool = False
-                    ) -> Iterator[dict]:
+    def _iterable_from_request(self, url, params):
+        # Shared by get_replays and get_groups
+        remaining = params["count"]
+        # return_length = True
+        while remaining > 0:
+            request_count = min(remaining, 200)
+            params["count"] = request_count
+            try:
+                d = self._request(url, "GET", params=params).json()
+            except HTTPError as e:
+                if e.response.status_code == 504:
+                    # Gateway Timeout, retry
+                    time.sleep(5)
+                    continue
+
+            batch = d["list"][:request_count]
+            yield from batch
+
+            if "next" not in d:
+                break
+
+            next_url = d["next"]
+            remaining -= len(batch)
+            params["after"] = parse_qs(urlparse(next_url).query)["after"][0]
+
+    def get_replays(
+            self,
+            *,
+            title: Optional[str] = None,
+            player_name: Optional[Union[str, List[str]]] = None,
+            player_id: Optional[Union[str, List[str]]] = None,
+            playlist: Optional[Union[AnyPlaylist, List[AnyPlaylist]]] = None,
+            season: Optional[Union[AnySeason, List[AnySeason]]] = None,
+            match_result: Optional[Union[AnyMatchResult, List[AnyMatchResult]]] = None,
+            min_rank: Optional[AnyRank] = None,
+            max_rank: Optional[AnyRank] = None,
+            pro: Optional[bool] = None,
+            uploader: Optional[str] = None,
+            group_id: Optional[Union[str, List[str]]] = None,
+            map_id: Optional[Union[AnyMap, List[AnyMap]]] = None,
+            created_before: Optional[Union[str, datetime]] = None,
+            created_after: Optional[Union[str, datetime]] = None,
+            replay_after: Optional[Union[str, datetime]] = None,
+            replay_before: Optional[Union[str, datetime]] = None,
+            count: int = 150,
+            sort_by: Optional[AnyReplaySortBy] = None,
+            sort_dir: AnySortDir = SortDir.DESCENDING,
+            deep: bool = False,
+            typed: Optional[bool] = None,
+    ) -> Iterator[Union[dict, ShallowReplay, DeepReplay]]:
         """
         This endpoint lets you filter and retrieve replays. The implementation returns an iterator.
 
@@ -184,42 +223,41 @@ class BallchasingApi:
         :param sort_by: sort replays according the selected field
         :param sort_dir: sort direction
         :param deep: whether to get full stats for each replay (will be much slower).
+        :param typed: whether to return a typed object (default is self.typed).
         :return: an iterator over the replays returned by the API.
         """
         url = f"{self.base_url}/replays"
         params = {"title": title, "player-name": player_name, "player-id": player_id, "playlist": playlist,
                   "season": season, "match-result": match_result, "min-rank": min_rank, "max-rank": max_rank,
                   "pro": pro, "uploader": uploader, "group": group_id, "map": map_id,
-                  "created-before": rfc3339(created_before), "created-after": rfc3339(created_after),
-                  "replay-date-after": rfc3339(replay_after), "replay-date-before": rfc3339(replay_before),
-                  "sort-by": sort_by, "sort-dir": sort_dir}
-        left = count
-        while left > 0:
-            request_count = min(left, 200)
-            params["count"] = request_count
-            d = self._request(url, "GET", params=params).json()
+                  "created-before": to_rfc3339(created_before), "created-after": to_rfc3339(created_after),
+                  "replay-date-after": to_rfc3339(replay_after), "replay-date-before": to_rfc3339(replay_before),
+                  "count": count, "sort-by": sort_by, "sort-dir": sort_dir}
 
-            batch = d["list"][:request_count]
-            if not deep:
-                yield from batch
-            else:
-                yield from (self.get_replay(r["id"]) for r in batch)
+        if typed is None:
+            typed = self.typed
 
-            if "next" not in d:
-                break
-
-            next_url = d["next"]
-            left -= len(batch)
-            params["after"] = parse_qs(urlparse(next_url).query)["after"][0]
+        iterator = self._iterable_from_request(url, params)
+        if deep:
+            iterator = (self.get_replay(r["id"], typed=typed) for r in iterator)
+        elif typed:
+            iterator = (ShallowReplay(**r) for r in iterator)
+        yield from iterator
 
-    def get_replay(self, replay_id: str) -> dict:
+    def get_replay(self, replay_id: str, *, typed: Optional[bool] = None) -> Union[dict, DeepReplay]:
         """
         Retrieve a given replay’s details and stats.
 
         :param replay_id: the replay id.
+        :param typed: whether to return a typed object (default is self.typed).
         :return: the result of the GET request.
         """
-        return self._request(f"/replays/{replay_id}", "GET").json()
+        result = self._request(f"/replays/{replay_id}", "GET").json()
+        if typed is None:
+            typed = self.typed
+        if typed:
+            result = DeepReplay(**result)
+        return result
 
     def patch_replay(self, replay_id: str, **params) -> None:
         """
@@ -230,18 +268,24 @@ class BallchasingApi:
         """
         self._request(f"/replays/{replay_id}", "PATCH", json=params)
 
-    def upload_replay(self,
-                      replay_file,
-                      visibility: Optional[AnyVisibility] = None,
-                      group: Optional[str] = None) -> dict:
+    def upload_replay(
+            self,
+            replay_file: Union[str, Path, BinaryIO],
+            *,
+            visibility: Optional[AnyVisibility] = None,
+            group: Optional[str] = None
+    ) -> dict:
         """
         Use this API to upload a replay file to ballchasing.com.
 
-        :param replay_file: replay file to upload.
+        :param replay_file: replay file to upload. Can be a file path (str or Path) or a file-like object.
         :param visibility: to set the visibility of the uploaded replay.
         :param group: to upload the replay to an existing group.
         :return: the result of the POST request.
         """
+        if isinstance(replay_file, (str, Path)):
+            with open(replay_file, "rb") as f:
+                return self.upload_replay(f, visibility=visibility, group=group)
         return self._request(f"/v2/upload", "POST", files={"file": replay_file},
                              params={"group": group, "visibility": visibility}).json()
 
@@ -254,16 +298,20 @@ class BallchasingApi:
         """
         self._request(f"/replays/{replay_id}", "DELETE")
 
-    def get_groups(self,
-                   name: Optional[str] = None,
-                   creator: Optional[str] = None,
-                   group: Optional[str] = None,
-                   created_before: Optional[Union[str, datetime]] = None,
-                   created_after: Optional[Union[str, datetime]] = None,
-                   count: int = 200,
-                   sort_by: AnyGroupSortBy = GroupSortBy.CREATED,
-                   sort_dir: AnySortDir = SortDir.DESCENDING
-                   ) -> Iterator[dict]:
+    def get_groups(
+            self,
+            *,
+            name: Optional[str] = None,
+            creator: Optional[str] = None,
+            group: Optional[str] = None,
+            created_before: Optional[Union[str, datetime]] = None,
+            created_after: Optional[Union[str, datetime]] = None,
+            count: int = 200,
+            sort_by: AnyGroupSortBy = GroupSortBy.CREATED,
+            sort_dir: AnySortDir = SortDir.DESCENDING,
+            deep: bool = False,
+            typed: bool = None,
+    ) -> Iterator[Union[dict, ShallowGroup, DeepGroup]]:
         """
         This endpoint lets you filter and retrieve replay groups.
 
@@ -279,34 +327,30 @@ class BallchasingApi:
                       past the limit of 200 set by the API
         :param sort_by: Sort groups according the selected field.
         :param sort_dir: Sort direction.
+        :param deep: whether to get full stats for each group (will be much slower).
+        :param typed: whether to return a typed object (default is self.typed).
         :return: an iterator over the groups returned by the API.
         """
         url = f"{self.base_url}/groups/"
-        params = {"name": name, "creator": creator, "group": group, "created-before": rfc3339(created_before),
-                  "created-after": rfc3339(created_after), "sort-by": sort_by, "sort-dir": sort_dir}
-
-        left = count
-        while left > 0:
-            request_count = min(left, 200)
-            params["count"] = request_count
-            d = self._request(url, "GET", params=params).json()
-
-            batch = d["list"][:request_count]
-            yield from batch
-
-            if "next" not in d:
-                break
-
-            next_url = d["next"]
-            left -= len(batch)
-            params["after"] = parse_qs(urlparse(next_url).query)["after"][0]
-
-    def create_group(self,
-                     name: str,
-                     player_identification: AnyPlayerIdentification,
-                     team_identification: AnyTeamIdentification,
-                     parent: Optional[str] = None
-                     ) -> dict:
+        params = {"name": name, "creator": creator, "group": group, "created-before": to_rfc3339(created_before),
+                  "created-after": to_rfc3339(created_after), "count": count, "sort-by": sort_by, "sort-dir": sort_dir}
+        iterator = self._iterable_from_request(url, params)
+        if typed is None:
+            typed = self.typed
+        if deep:
+            iterator = (self.get_group(g["id"], typed=typed) for g in iterator)
+        elif typed:
+            iterator = (ShallowGroup(**g) for g in iterator)
+        yield from iterator
+
+    def create_group(
+            self,
+            *,
+            name: str,
+            player_identification: AnyPlayerIdentification,
+            team_identification: AnyTeamIdentification,
+            parent: Optional[str] = None
+    ) -> dict:
         """
         Use this API to create a new replay group.
 
@@ -326,14 +370,25 @@ class BallchasingApi:
                 "team_identification": team_identification, "parent": parent}
         return self._request(f"/groups", "POST", json=json).json()
 
-    def get_group(self, group_id: str) -> dict:
+    def get_group(
+            self,
+            group_id: str,
+            *,
+            typed: Optional[bool] = None
+    ) -> Union[dict, DeepGroup]:
         """
         This endpoint retrieves a specific replay group info and stats given its id.
 
         :param group_id: the group id.
+        :param typed: whether to return a typed object (default is self.typed).
         :return: the group info with stats.
         """
-        return self._request(f"/groups/{group_id}", "GET").json()
+        result = self._request(f"/groups/{group_id}", "GET").json()
+        if typed is None:
+            typed = self.typed
+        if typed:
+            result = DeepGroup(**result)
+        return result
 
     def patch_group(self, group_id: str, **params) -> None:
         """
@@ -353,46 +408,79 @@ class BallchasingApi:
         """
         self._request(f"/groups/{group_id}", "DELETE")
 
-    def get_group_replays(self, group_id: str, deep: bool = False) -> Iterator[dict]:
+    def get_group_replays(
+            self,
+            group: Union[str, dict, BasicGroup],
+            *,
+            deep: bool = False,
+            typed: Optional[bool] = None
+    ) -> Iterator[Union[dict, ShallowReplay, DeepReplay]]:
         """
         Finds all replays in a group, including child groups.
 
-        :param group_id: the base group id.
+        :param group: the base group id, group dict, or BaseGroup object.
         :param deep: whether or not to get full stats for each replay (will be much slower).
+        :param typed: whether to return a typed object (default is self.typed).
         :return: an iterator over all the replays in the group.
         """
-        child_groups = self.get_groups(group=group_id)
-        for child in child_groups:
-            for replay in self.get_group_replays(child["id"], deep):
-                yield replay
-        for replay in self.get_replays(group_id=group_id, deep=deep):
+        for path in self.get_group_tree(group, deep=deep, typed=typed):
+            group, replay = path
             yield replay
 
-    def download_replay(self, replay_id: str, folder: str):
+    def get_group_tree(
+            self,
+            group: Union[str, dict, BaseGroup],
+            *,
+            deep: bool = False,
+            typed: Optional[bool] = None
+    ):
+        """
+        Finds all replays in a group, and includes the groups leading up to the replays.
+        :param group: the group id or a group dict.
+        :param deep: whether to get full stats for each replay and group (will be much slower).
+        :param typed: whether to return a typed object (default is self.typed).
+        """
+        if isinstance(group, str):
+            group = self.get_group(group)
+        if isinstance(group, BasicGroup):
+            group_id = group.id
+        else:
+            group_id = group["id"]
+        child_groups = self.get_groups(group=group_id, typed=typed)
+        for child in child_groups:
+            for path in self.get_group_tree(child, deep=deep, typed=typed):
+                yield group_id, *path
+        for replay in self.get_replays(group_id=group_id, deep=deep, typed=typed):
+            yield group_id, replay
+
+    def download_replay(self, replay_id: str, path: str):
         """
         Download a replay file.
 
         :param replay_id: the replay id.
-        :param folder: the folder to download into.
+        :param path: the path to download the replay to. Can be a file path or a directory.
         """
         r = self._request(f"/replays/{replay_id}/file", "GET")
-        with open(f"{folder}/{replay_id}.replay", "wb") as f:
-            for ch in r:
-                f.write(ch)
+        if os.path.isdir(path):
+            # If path is a directory, use the replay id as the filename
+            filename = f"{replay_id}.replay"
+            path = os.path.join(path, filename)
+        with open(path, "wb") as f:
+            f.write(r.content)
 
-    def download_group(self, group_id: str, folder: str, recursive=True):
+    def download_group(self, group_id: str, folder: str, *, keep_tree_structure=True):
         """
         Download an entire group.
 
         :param group_id: the base group id.
         :param folder: the folder in which to create the group folder.
-        :param recursive: whether to create new folders for child groups.
+        :param keep_tree_structure: whether to create new folders for child groups.
         """
         folder = os.path.join(folder, group_id)
-        if recursive:
+        if keep_tree_structure:
             os.makedirs(folder, exist_ok=True)
             for child_group in self.get_groups(group=group_id):
-                self.download_group(child_group["id"], folder, True)
+                self.download_group(child_group["id"], folder, keep_tree_structure=True)
             for replay in self.get_replays(group_id=group_id):
                 self.download_replay(replay["id"], folder)
         else:
@@ -415,9 +503,9 @@ class BallchasingApi:
         """
 
         if isinstance(replay, str):
-            replay = self.get_replay(replay)
+            replay = self.get_replay(replay, typed=False)
         elif isinstance(replay, dict) and "title" not in replay:
-            replay = self.get_replay(replay["id"])
+            replay = self.get_replay(replay["id"], typed=False)
 
         stats = parse_replay_stats(replay)
         return stats