Allow the user to set values of query params for studies

berserk/clients/studies.py

@@ -11,32 +11,98 @@
 class Studies(BaseClient):
     """Study chess the Lichess way."""
 
-    def export_chapter(self, study_id: str, chapter_id: str) -> str:
+    def export_chapter(
+        self,
+        study_id: str,
+        chapter_id: str,
+        clocks: bool = True,
+        comments: bool = True,
+        variations: bool = True,
+        source: bool = False,
+        orientation: bool = False,
+    ) -> str:
         """Export one chapter of a study.
 
+        :param study_id: study id
+        :param chapter_id: chapter id
+        :param clocks: include any clock comments in the PGN moves
+        :param comments: include any analysis and annotator comments in the PGN moves
+        :param variations: include any variations in the PGN moves
+        :param source: include a `Source` PGN tag containing the chapter URL
+        :param orientation: include an `Orientation` PGN tag containing the chapter's board orientation
         :return: chapter PGN
         """
         path = f"/api/study/{study_id}/{chapter_id}.pgn"
-        return self._r.get(path, fmt=PGN)
+        params = {
+            "clocks": clocks,
+            "comments": comments,
+            "variations": variations,
+            "source": source,
+            "orientation": orientation,
+        }
+        return self._r.get(path, fmt=PGN, params=params)
 
-    def export(self, study_id: str) -> Iterator[str]:
+    def export(
+        self,
+        study_id: str,
+        clocks: bool = True,
+        comments: bool = True,
+        variations: bool = True,
+        source: bool = False,
+        orientation: bool = False,
+    ) -> Iterator[str]:
         """Export all chapters of a study.
 
-        :return: iterator over all chapters as PGN
+        :param study_id: study id
+        :param clocks: include any clock comments in the PGN moves
+        :param comments: include any analysis and annotator comments in the PGN moves
+        :param variations: include any variations in the PGN moves
+        :param source: for each chapter, include a `Source` PGN tag containing the chapter URL
+        :param orientation: for each chapter, include an `Orientation` PGN tag containing the chapter's board orientation
+        :return: iterator over all chapters as PGNs
         """
         path = f"/api/study/{study_id}.pgn"
-        yield from self._r.get(path, fmt=PGN, stream=True)
+        params = {
+            "clocks": clocks,
+            "comments": comments,
+            "variations": variations,
+            "source": source,
+            "orientation": orientation,
+        }
+        yield from self._r.get(path, fmt=PGN, stream=True, params=params)
 
-    def export_by_username(self, username: str) -> Iterator[str]:
+    def export_by_username(
+        self,
+        username: str,
+        clocks: bool = True,
+        comments: bool = True,
+        variations: bool = True,
+        source: bool = False,
+        orientation: bool = False,
+    ) -> Iterator[str]:
         """Export all chapters of all studies of a user in PGN format.
 
         If authenticated, then all public, unlisted, and private studies are included.
 
         If not, only public (non-unlisted) studies are included.
 
-        return:iterator over all chapters as PGN"""
+        :param username: the user whose studies will be exported
+        :param clocks: include any clock comments in the PGN moves
+        :param comments: include any analysis and annotator comments in the PGN moves
+        :param variations: include any variations in the PGN moves
+        :param source: for each chapter, include a `Source` PGN tag containing the chapter URL
+        :param orientation: for each chapter, include an `Orientation` PGN tag containing the chapter's board orientation
+        :return: iterator over all chapters as PGNs
+        """
         path = f"/study/by/{username}/export.pgn"
-        yield from self._r.get(path, fmt=PGN, stream=True)
+        params = {
+            "clocks": clocks,
+            "comments": comments,
+            "variations": variations,
+            "source": source,
+            "orientation": orientation,
+        }
+        yield from self._r.get(path, fmt=PGN, stream=True, params=params)
 
     def import_pgn(
         self,
@@ -46,7 +112,7 @@ def import_pgn(
         orientation: Color = "white",
         variant: Variant = "standard",
     ) -> List[ChapterIdName]:
-        """Imports arbitrary PGN into an existing study.
+        """Imports an arbitrary PGN into an existing study.
         Creates a new chapter in the study.
 
         If the PGN contains multiple games (separated by 2 or more newlines) then multiple chapters will be created within the study.