Add polls

disgoorg · Mar 21, 2024 · f842904 · f842904
1 parent 0f6192b
commit f842904
Show file tree

Hide file tree

Showing 14 changed files with 175 additions and 0 deletions.
diff --git a/discord/message.go b/discord/message.go
@@ -114,6 +114,7 @@ type Message struct {
 	Position             *int                  `json:"position,omitempty"`
 	RoleSubscriptionData *RoleSubscriptionData `json:"role_subscription_data,omitempty"`
 	Resolved             *ResolvedData         `json:"resolved,omitempty"`
+	Poll                 *Poll                 `json:"poll,omitempty"`
 }
 
 func (m *Message) UnmarshalJSON(data []byte) error {

diff --git a/discord/message_create.go b/discord/message_create.go
@@ -18,6 +18,7 @@ type MessageCreate struct {
 	MessageReference *MessageReference    `json:"message_reference,omitempty"`
 	Flags            MessageFlags         `json:"flags,omitempty"`
 	EnforceNonce     bool                 `json:"enforce_nonce,omitempty"`
+	Poll             Poll                 `json:"poll,omitempty"`
 }
 
 func (MessageCreate) interactionCallbackData() {}

diff --git a/discord/message_create_builder.go b/discord/message_create_builder.go
@@ -251,6 +251,12 @@ func (b *MessageCreateBuilder) SetSuppressEmbeds(suppressEmbeds bool) *MessageCr
 	return b
 }
 
+// SetPoll sets the Poll of the Message
+func (b *MessageCreateBuilder) SetPoll(poll Poll) *MessageCreateBuilder {
+	b.Poll = poll
+	return b
+}
+
 // Build builds the MessageCreateBuilder to a MessageCreate struct
 func (b *MessageCreateBuilder) Build() MessageCreate {
 	return b.MessageCreate

diff --git a/discord/poll.go b/discord/poll.go
@@ -0,0 +1,43 @@
+package discord
+
+import (
+	"time"
+
+	"github.com/disgoorg/snowflake/v2"
+)
+
+type Poll struct {
+	Question         PollMedia      `json:"question"`
+	Answers          []PollAnswer   `json:"answers"`
+	Expiry           time.Time      `json:"expiry"`
+	AllowMultiselect bool           `json:"allow_multiselect"`
+	LayoutType       PollLayoutType `json:"layout_type"`
+	Results          []PollResults  `json:"results"`
+}
+
+type PollMedia struct {
+	Text  *string       `json:"text"`
+	Emoji *PartialEmoji `json:"emoji"`
+}
+
+type PollAnswer struct {
+	AnswerID  *snowflake.ID `json:"answer_id"`
+	PollMedia PollMedia     `json:"poll_media"`
+}
+
+type PollResults struct {
+	IsFinalized  bool              `json:"is_finalized"`
+	AnswerCounts []PollAnswerCount `json:"answer_counts"`
+}
+
+type PollAnswerCount struct {
+	ID      snowflake.ID `json:"id"`
+	Count   int          `json:"count"`
+	MeVoted bool         `json:"me_voted"`
+}
+
+type PollLayoutType int
+
+const (
+	PollLayoutTypeDefault PollLayoutType = iota + 1
+)
diff --git a/discord/webhook_message_create.go b/discord/webhook_message_create.go
@@ -15,6 +15,7 @@ type WebhookMessageCreate struct {
 	Flags           MessageFlags         `json:"flags,omitempty"`
 	ThreadName      string               `json:"thread_name,omitempty"`
 	AppliedTags     []snowflake.ID       `json:"applied_tags,omitempty"`
+	Poll            Poll                 `json:"poll,omitempty"`
 }
 
 // ToBody returns the MessageCreate ready for body

diff --git a/discord/webhook_message_create_builder.go b/discord/webhook_message_create_builder.go
@@ -210,6 +210,12 @@ func (b *WebhookMessageCreateBuilder) SetThreadName(threadName string) *WebhookM
 	return b
 }
 
+// SetPoll sets the Poll of the webhook Message
+func (b *WebhookMessageCreateBuilder) SetPoll(poll Poll) *WebhookMessageCreateBuilder {
+	b.Poll = poll
+	return b
+}
+
 // Build builds the WebhookMessageCreateBuilder to a MessageCreate struct
 func (b *WebhookMessageCreateBuilder) Build() WebhookMessageCreate {
 	b.WebhookMessageCreate.Components = b.Components

diff --git a/events/listener_adapter.go b/events/listener_adapter.go
@@ -147,6 +147,10 @@ type ListenerAdapter struct {
 	OnMessageUpdate func(event *MessageUpdate)
 	OnMessageDelete func(event *MessageDelete)
 
+	// Message Poll Events
+	OnMessagePollVoteAdd    func(event *MessagePollVoteAdd)
+	OnMessagePollVoteRemove func(event *MessagePollVoteRemove)
+
 	// Message Reaction Events
 	OnMessageReactionAdd         func(event *MessageReactionAdd)
 	OnMessageReactionRemove      func(event *MessageReactionRemove)
@@ -574,6 +578,16 @@ func (l *ListenerAdapter) OnEvent(event bot.Event) {
 			listener(e)
 		}
 
+	// Message Poll Events
+	case *MessagePollVoteAdd:
+		if listener := l.OnMessagePollVoteAdd; listener != nil {
+			listener(e)
+		}
+	case *MessagePollVoteRemove:
+		if listener := l.OnMessagePollVoteRemove; listener != nil {
+			listener(e)
+		}
+
 	// Message Reaction Events
 	case *MessageReactionAdd:
 		if listener := l.OnMessageReactionAdd; listener != nil {

diff --git a/events/message_events.go b/events/message_events.go
@@ -43,3 +43,23 @@ type MessageUpdate struct {
 type MessageDelete struct {
 	*GenericMessage
 }
+
+// GenericMessagePoll is a generic poll event
+type GenericMessagePoll struct {
+	*GenericEvent
+	UserID    snowflake.ID
+	ChannelID snowflake.ID
+	MessageID snowflake.ID
+	GuildID   *snowflake.ID
+	AnswerID  snowflake.ID
+}
+
+// MessagePollVoteAdd indicates that a discord.User voted on a discord.Poll
+type MessagePollVoteAdd struct {
+	*GenericMessagePoll
+}
+
+// MessagePollVoteRemove indicates that a discord.User removed their vote on a discord.Poll
+type MessagePollVoteRemove struct {
+	*GenericMessagePoll
+}
diff --git a/gateway/gateway_event_type.go b/gateway/gateway_event_type.go
@@ -59,6 +59,8 @@ const (
 	EventTypeMessageUpdate                       EventType = "MESSAGE_UPDATE"
 	EventTypeMessageDelete                       EventType = "MESSAGE_DELETE"
 	EventTypeMessageDeleteBulk                   EventType = "MESSAGE_DELETE_BULK"
+	EventTypeMessagePollVoteAdd                  EventType = "MESSAGE_POLL_VOTE_ADD"
+	EventTypeMessagePollVoteRemove               EventType = "MESSAGE_POLL_VOTE_REMOVE"
 	EventTypeMessageReactionAdd                  EventType = "MESSAGE_REACTION_ADD"
 	EventTypeMessageReactionRemove               EventType = "MESSAGE_REACTION_REMOVE"
 	EventTypeMessageReactionRemoveAll            EventType = "MESSAGE_REACTION_REMOVE_ALL"

diff --git a/gateway/gateway_events.go b/gateway/gateway_events.go
@@ -507,6 +507,28 @@ type EventMessageDeleteBulk struct {
 func (EventMessageDeleteBulk) messageData() {}
 func (EventMessageDeleteBulk) eventData()   {}
 
+type EventMessagePollVoteAdd struct {
+	UserID    snowflake.ID  `json:"user_id"`
+	ChannelID snowflake.ID  `json:"channel_id"`
+	MessageID snowflake.ID  `json:"message_id"`
+	GuildID   *snowflake.ID `json:"guild_id"`
+	AnswerID  snowflake.ID  `json:"answer_id"`
+}
+
+func (EventMessagePollVoteAdd) messageData() {}
+func (EventMessagePollVoteAdd) eventData()   {}
+
+type EventMessagePollVoteRemove struct {
+	UserID    snowflake.ID  `json:"user_id"`
+	ChannelID snowflake.ID  `json:"channel_id"`
+	MessageID snowflake.ID  `json:"message_id"`
+	GuildID   *snowflake.ID `json:"guild_id"`
+	AnswerID  snowflake.ID  `json:"answer_id"`
+}
+
+func (EventMessagePollVoteRemove) messageData() {}
+func (EventMessagePollVoteRemove) eventData()   {}
+
 type EventPresenceUpdate struct {
 	discord.Presence
 }

diff --git a/handlers/all_handlers.go b/handlers/all_handlers.go
@@ -101,6 +101,9 @@ var allEventHandlers = []bot.GatewayEventHandler{
 	bot.NewGatewayEventHandler(gateway.EventTypeMessageDelete, gatewayHandlerMessageDelete),
 	bot.NewGatewayEventHandler(gateway.EventTypeMessageDeleteBulk, gatewayHandlerMessageDeleteBulk),
 
+	bot.NewGatewayEventHandler(gateway.EventTypeMessagePollVoteAdd, gatewayHandlerMessagePollVoteAdd),
+	bot.NewGatewayEventHandler(gateway.EventTypeMessagePollVoteRemove, gatewayHandlerMessagePollVoteRemove),
+
 	bot.NewGatewayEventHandler(gateway.EventTypeMessageReactionAdd, gatewayHandlerMessageReactionAdd),
 	bot.NewGatewayEventHandler(gateway.EventTypeMessageReactionRemove, gatewayHandlerMessageReactionRemove),
 	bot.NewGatewayEventHandler(gateway.EventTypeMessageReactionRemoveAll, gatewayHandlerMessageReactionRemoveAll),

diff --git a/handlers/message_poll_handler.go b/handlers/message_poll_handler.go
@@ -0,0 +1,33 @@
+package handlers
+
+import (
+	"github.com/disgoorg/disgo/bot"
+	"github.com/disgoorg/disgo/events"
+	"github.com/disgoorg/disgo/gateway"
+)
+
+func gatewayHandlerMessagePollVoteAdd(client bot.Client, sequenceNumber int, shardID int, event gateway.EventMessagePollVoteAdd) {
+	client.EventManager().DispatchEvent(&events.MessagePollVoteAdd{
+		GenericMessagePoll: &events.GenericMessagePoll{
+			GenericEvent: events.NewGenericEvent(client, sequenceNumber, shardID),
+			UserID:       event.UserID,
+			ChannelID:    event.ChannelID,
+			MessageID:    event.MessageID,
+			GuildID:      event.GuildID,
+			AnswerID:     event.AnswerID,
+		},
+	})
+}
+
+func gatewayHandlerMessagePollVoteRemove(client bot.Client, sequenceNumber int, shardID int, event gateway.EventMessagePollVoteRemove) {
+	client.EventManager().DispatchEvent(&events.MessagePollVoteRemove{
+		GenericMessagePoll: &events.GenericMessagePoll{
+			GenericEvent: events.NewGenericEvent(client, sequenceNumber, shardID),
+			UserID:       event.UserID,
+			ChannelID:    event.ChannelID,
+			MessageID:    event.MessageID,
+			GuildID:      event.GuildID,
+			AnswerID:     event.AnswerID,
+		},
+	})
+}
diff --git a/rest/channels.go b/rest/channels.go
@@ -47,6 +47,9 @@ type Channels interface {
 	PinMessage(channelID snowflake.ID, messageID snowflake.ID, opts ...RequestOpt) error
 	UnpinMessage(channelID snowflake.ID, messageID snowflake.ID, opts ...RequestOpt) error
 	Follow(channelID snowflake.ID, targetChannelID snowflake.ID, opts ...RequestOpt) (*discord.FollowedChannel, error)
+
+	GetPollAnswerVotes(channelID snowflake.ID, messageID snowflake.ID, answerID snowflake.ID, after snowflake.ID, limit int, opts ...RequestOpt) ([]discord.User, error)
+	ExpirePoll(channelID snowflake.ID, messageID snowflake.ID, opts ...RequestOpt) (*discord.Message, error)
 }
 
 type channelImpl struct {
@@ -222,3 +225,20 @@ func (s *channelImpl) Follow(channelID snowflake.ID, targetChannelID snowflake.I
 	err = s.client.Do(FollowChannel.Compile(nil, channelID), discord.FollowChannel{ChannelID: targetChannelID}, &followedChannel, opts...)
 	return
 }
+
+func (s *channelImpl) GetPollAnswerVotes(channelID snowflake.ID, messageID snowflake.ID, answerID snowflake.ID, after snowflake.ID, limit int, opts ...RequestOpt) (users []discord.User, err error) {
+	values := discord.QueryValues{}
+	if after != 0 {
+		values["after"] = after
+	}
+	if limit != 0 {
+		values["limit"] = limit
+	}
+	err = s.client.Do(GetPollAnswerVotes.Compile(values, channelID, messageID, answerID), nil, &users, opts...)
+	return
+}
+
+func (s *channelImpl) ExpirePoll(channelID snowflake.ID, messageID snowflake.ID, opts ...RequestOpt) (message *discord.Message, err error) {
+	err = s.client.Do(ExpirePoll.Compile(nil, channelID, messageID), nil, &message, opts...)
+	return
+}
diff --git a/rest/rest_endpoints.go b/rest/rest_endpoints.go
@@ -170,6 +170,9 @@ var (
 
 	SendTyping    = NewEndpoint(http.MethodPost, "/channels/{channel.id}/typing")
 	FollowChannel = NewEndpoint(http.MethodPost, "/channels/{channel.id}/followers")
+
+	GetPollAnswerVotes = NewEndpoint(http.MethodGet, "/channels/{channel.id}/polls/{message.id}/answers/{answer.id}")
+	ExpirePoll         = NewEndpoint(http.MethodPost, "/channels/{channel.id}/poll/{message.id}/expire")
 )
 
 // Threads