Add support for creating, getting, and resolving predictions (#207)

Resolves #178

Author: raechelmaelstrom

twitchio/http.py

@@ -520,6 +520,63 @@ async def update_reward_redemption_status(
             )
         )
 
+    async def get_predictions(
+        self,
+        token: str,
+        broadcaster_id: int,
+        prediction_id: str = None,
+    ):
+        params = [("broadcaster_id", str(broadcaster_id))]
+
+        if prediction_id:
+            params.extend(("prediction_id", prediction_id))
+
+        return await self.request(Route("GET", "predictions", query=params, token=token), paginate=False)
+
+    async def patch_prediction(
+            self,
+            token: str,
+            broadcaster_id: int,
+            prediction_id: str,
+            status: str,
+            winning_outcome_id: str = None):
+        body = {
+            "broadcaster_id": str(broadcaster_id),
+            "id": prediction_id,
+            "status": status,
+        }
+
+        if status == "RESOLVED":
+            body["winning_outcome_id"] = winning_outcome_id
+
+        return await self.request(
+            Route(
+                "PATCH",
+                "predictions",
+                body=body,
+                token=token,
+            )
+        )
+
+    async def post_prediction(self, token: str, broadcaster_id: int, title: str, blue_outcome: str, pink_outcome: str, prediction_window: int):
+        body = {
+            "broadcaster_id": broadcaster_id,
+            "title": title,
+            "prediction_window": prediction_window,
+            "outcomes": [
+                {
+                    "title": blue_outcome,
+                },
+                {
+                    "title": pink_outcome,
+                }
+            ],
+        }
+        return await self.request(
+            Route("POST", "predictions", body=body, token=token),
+            paginate=False,
+        )
+
     async def post_create_clip(self, token: str, broadcaster_id: int, has_delay=False):
         return await self.request(
             Route("POST", "clips", query=[("broadcaster_id", broadcaster_id), ("has_delay", has_delay)], token=token),

twitchio/models.py

@@ -55,6 +55,7 @@
     "Video",
     "Tag",
     "WebhookSubscription",
+    "Prediction",
 )
 
 
@@ -581,3 +582,62 @@ def __init__(self, http: "TwitchHTTP", data: dict):
 
     def __repr__(self):
         return f"<ChannelInfo user={self.user} game_id={self.game_id} game_name={self.game_name} title={self.title} language={self.language} delay={self.delay}>"
+
+
+class Prediction:
+
+    __slots__ = ("user", "prediction_id", "title", "winning_outcome_id", "outcomes", "prediction_window", "prediction_status", "created_at", "ended_at", "locked_at")
+
+    def __init__(self, http: "TwitchHTTP", data: dict):
+        self.user = PartialUser(http, data["broadcaster_id"], data["broadcaster_name"])
+        self.prediction_id: str = data["id"]
+        self.title: str = data["title"]
+        self.winning_outcome_id: str = data["winning_outcome_id"]
+        self.outcomes: List[PredictionOutcome] = [PredictionOutcome(http, x) for x in data["outcomes"]]
+        self.prediction_window: int = data["prediction_window"]
+        self.prediction_status: str = data["status"]
+        self.created_at = self._parse_time(data, "created_at")
+        self.ended_at = self._parse_time(data, "ended_at")
+        self.locked_at = self._parse_time(data, "locked_at")
+
+    def _parse_time(self, data, field) -> Optional["Datetime"]:
+        if field not in data or data[field] == None:
+            return None
+
+        time = data[field].split(".")[0]
+        return datetime.datetime.fromisoformat(time)
+
+    def __repr__(self):
+        return f"<Prediction user={self.user} prediction_id={self.prediction_id} winning_outcome_id={self.winning_outcome_id} title={self.title}>"
+
+class Predictor:
+
+    __slots__ = ("outcome_id", "title", "channel_points", "color")
+
+    def __init__(self, http: "TwitchHTTP", data: dict):
+        self.channel_points_used: int = data["channel_points_used"]
+        self.channel_points_won: int = data["channel_points_won"]
+        self.user = PartialUser(http, data["user"]["id"], data["user"]["name"])
+
+class PredictionOutcome:
+
+    __slots__ = ("outcome_id", "title", "channel_points", "color", "users", "top_predictors")
+
+    def __init__(self, http: "TwitchHTTP", data: dict):
+        self.outcome_id: str = data["id"]
+        self.title: str = data["title"]
+        self.channel_points: int = data["channel_points"]
+        self.color: str = data["color"]
+        self.users: int = data["users"]
+        if data["top_predictors"]:
+            self.top_predictors: List[Predictor] = [Predictor(http, x) for x in data["top_predictors"]]
+        else:
+            self.top_predictors: List[Predictor] = None
+
+    @property
+    def colour(self) -> str:
+        """The colour of the prediction. Alias to color."""
+        return self.color
+
+    def __repr__(self):
+        return f"<PredictionOutcome outcome_id={self.outcome_id} title={self.title} channel_points={self.channel_points} color={self.color}>"

twitchio/user.py

@@ -34,7 +34,7 @@
 if TYPE_CHECKING:
     from .http import TwitchHTTP
     from .channel import Channel
-    from .models import BitsLeaderboard, Clip, ExtensionBuilder, Tag, FollowEvent
+    from .models import BitsLeaderboard, Clip, ExtensionBuilder, Tag, FollowEvent, Prediction
 
 
 __all__ = (
@@ -609,6 +609,87 @@ async def fetch_videos(self, period="all", sort="time", type="all", language=Non
         data = await self._http.get_videos(user_id=str(self.id), period=period, sort=sort, type=type, language=language)
         return [Video(self._http, x, self) for x in data]
 
+    async def end_prediction(self, token: str, prediction_id: str, status: str, winning_outcome_id: str = None) -> "Prediction":
+        """|coro|
+        End a prediction with an outcome.
+
+        Parameters
+        -----------
+        token: :class:`str`
+            An oauth token with the channel:manage:predictions scope
+        prediction_id: :class:`str`
+            ID of the prediction to end.
+
+        Returns
+        --------
+            :class:`twitchio.Prediction`
+        """
+        from .models import Prediction
+
+        data = await self._http.patch_prediction(
+                token,
+                broadcaster_id=str(self.id),
+                prediction_id=prediction_id,
+                status=status,
+                winning_outcome_id=winning_outcome_id)
+        return Prediction(self._http, data[0])
+
+    async def get_predictions(self, token: str, prediction_id: str = None) -> List["Prediction"]:
+        """|coro|
+        Gets information on a prediction or the list of predictions
+        if none is provided.
+
+        Parameters
+        -----------
+        token: :class:`str`
+            An oauth token with the channel:manage:predictions scope
+        prediction_id: :class:`str`
+            ID of the prediction to receive information about.
+
+        Returns
+        --------
+            :class:`twitchio.Prediction`
+        """
+        from .models import Prediction
+
+        data = await self._http.get_predictions(
+                token,
+                broadcaster_id=str(self.id),
+                prediction_id=prediction_id)
+        return [Prediction(self._http, d) for d in data]
+
+    async def create_prediction(self, token: str, title: str, blue_outcome: str, pink_outcome: str, prediction_window: int) -> "Prediction":
+        """|coro|
+        Creates a prediction for the channel.
+
+        Parameters
+        -----------
+        token: :class:`str`
+            An oauth token with the channel:manage:predictions scope
+        title: :class:`str`
+            Title for the prediction (max of 45 characters)
+        blue_outcome: :class:`str`
+            Text for the first outcome people can vote for. (max 25 characters)
+        pink_outcome: :class:`str`
+            Text for the second outcome people can vote for. (max 25 characters)
+        prediction_window: :class:`int`
+            Total duration for the prediction (in seconds)
+
+        Returns
+        --------
+            :class:`twitchio.Prediction`
+        """
+        from .models import Prediction
+
+        data = await self._http.post_prediction(
+                token,
+                broadcaster_id=str(self.id),
+                title=title,
+                blue_outcome=blue_outcome,
+                pink_outcome=pink_outcome,
+                prediction_window=prediction_window)
+        return Prediction(self._http, data[0])
+
 
 class BitLeaderboardUser(PartialUser):