diff --git a/berserk/clients/studies.py b/berserk/clients/studies.py index c9bac0e..d92cec4 100644 --- a/berserk/clients/studies.py +++ b/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.