diff --git a/python_sponsorblock/__init__.py b/python_sponsorblock/__init__.py index 114ad1d..4e829d3 100644 --- a/python_sponsorblock/__init__.py +++ b/python_sponsorblock/__init__.py @@ -7,7 +7,7 @@ import json from functools import wraps from .exceptions import SponsorBlockError, SponsorBlockIdNotFoundError, ReturnDefault -from .constants import Segment +from .constants import Segment, Category def error_handling(default: Any) -> Callable: @@ -99,14 +99,25 @@ class SponsorBlock: return data @error_handling(default=[]) - def get_segments(self, video: str) -> List[Segment]: + def get_segments(self, video: str, categories: List[Category] = None) -> List[Segment]: + """ + Retrieves the skip segments for a given video. + + Args: + video (str): The video identifier. + categories (List[Category], optional): A list of categories to filter the skip segments. Defaults to all categories. + + Returns: + List[Segment]: A list of skip segments for the given video. + """ video_id = self._get_video_id(video) + categories = categories or [c for c in Category] # build query parameters query = { - "videoID": video_id + "videoID": video_id, + "categories": json.dumps([c.value for c in categories]) } - print(query) r = self._request(method="GET", endpoint="/api/skipSegments?" + urlencode(query)) return [constants.Segment(**d) for d in r]