Browse Source

Implement state tracking.

Currently maintains a list of Guilds (Members/Channels) and PrivateChannels.
Chris Rhodes 9 years ago
parent
commit
4e23e5fc35
4 changed files with 257 additions and 47 deletions
  1. 3 1
      discord.go
  2. 195 0
      state.go
  3. 12 10
      structs.go
  4. 47 36
      wsapi.go

+ 3 - 1
discord.go

@@ -61,7 +61,9 @@ type voiceUDP struct {
 func New(args ...interface{}) (s *Session, err error) {
 
 	// Create an empty Session interface.
-	s = &Session{}
+	s = &Session{
+		State: NewState(),
+	}
 
 	// If no arguments are passed return the empty Session interface.
 	// Later I will add default values, if appropriate.

+ 195 - 0
state.go

@@ -0,0 +1,195 @@
+package discordgo
+
+import "errors"
+
+// NewState creates an empty state.
+func NewState() *State {
+	return &State{
+		Ready: Ready{
+			PrivateChannels: []Channel{},
+			Guilds:          []Guild{},
+		},
+	}
+}
+
+// OnReady takes a Ready event and updates all internal state.
+func (s *State) OnReady(r *Ready) {
+	s.Ready = *r
+}
+
+// AddGuild adds a guild to the current world state, or
+// updates it if it already exists.
+func (s *State) AddGuild(guild *Guild) {
+	for _, g := range s.Guilds {
+		if g.ID == guild.ID {
+			// This could be a little faster ;)
+			for _, m := range guild.Members {
+				s.AddMember(&m)
+			}
+			for _, c := range guild.Channels {
+				s.AddChannel(&c)
+			}
+			return
+		}
+	}
+	s.Guilds = append(s.Guilds, *guild)
+}
+
+// RemoveGuild removes a guild from current world state.
+func (s *State) RemoveGuild(guild *Guild) error {
+	for i, g := range s.Guilds {
+		if g.ID == guild.ID {
+			s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...)
+			return nil
+		}
+	}
+	return errors.New("Guild not found.")
+}
+
+// GetGuildByID gets a guild by ID.
+// Useful for querying if @me is in a guild:
+//     _, err := discordgo.Session.State.GetGuildById(guildID)
+//     isInGuild := err == nil
+func (s *State) GetGuildByID(guildID string) (*Guild, error) {
+	for _, g := range s.Guilds {
+		if g.ID == guildID {
+			return &g, nil
+		}
+	}
+	return nil, errors.New("Guild not found.")
+}
+
+// TODO: Consider moving Guild state update methods onto *Guild.
+
+// AddMember adds a member to the current world state, or
+// updates it if it already exists.
+func (s *State) AddMember(member *Member) error {
+	guild, err := s.GetGuildByID(member.GuildID)
+	if err != nil {
+		return err
+	}
+
+	for i, m := range guild.Members {
+		if m.User.ID == member.User.ID {
+			guild.Members[i] = *member
+			return nil
+		}
+	}
+
+	guild.Members = append(guild.Members, *member)
+	return nil
+}
+
+// RemoveMember removes a member from current world state.
+func (s *State) RemoveMember(member *Member) error {
+	guild, err := s.GetGuildByID(member.GuildID)
+	if err != nil {
+		return err
+	}
+
+	for i, m := range guild.Members {
+		if m.User.ID == member.User.ID {
+			guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
+			return nil
+		}
+	}
+	return errors.New("Member not found.")
+}
+
+// GetMemberByID gets a member by ID from a guild.
+func (s *State) GetMemberByID(guildID string, userID string) (*Member, error) {
+	guild, err := s.GetGuildByID(guildID)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, m := range guild.Members {
+		if m.User.ID == userID {
+			return &m, nil
+		}
+	}
+	return nil, errors.New("Member not found.")
+}
+
+// AddChannel adds a guild to the current world state, or
+// updates it if it already exists.
+// Channels may exist either as PrivateChannels or inside
+// a guild.
+func (s *State) AddChannel(channel *Channel) error {
+	if channel.IsPrivate {
+		for i, c := range s.PrivateChannels {
+			if c.ID == channel.ID {
+				s.PrivateChannels[i] = *channel
+				return nil
+			}
+		}
+
+		s.PrivateChannels = append(s.PrivateChannels, *channel)
+	} else {
+		guild, err := s.GetGuildByID(channel.GuildID)
+		if err != nil {
+			return err
+		}
+
+		for i, c := range guild.Channels {
+			if c.ID == channel.ID {
+				guild.Channels[i] = *channel
+				return nil
+			}
+		}
+
+		guild.Channels = append(guild.Channels, *channel)
+	}
+	return nil
+}
+
+// RemoveChannel removes a channel from current world state.
+func (s *State) RemoveChannel(channel *Channel) error {
+	if channel.IsPrivate {
+		for i, c := range s.PrivateChannels {
+			if c.ID == channel.ID {
+				s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...)
+				return nil
+			}
+		}
+	} else {
+		guild, err := s.GetGuildByID(channel.GuildID)
+		if err != nil {
+			return err
+		}
+
+		for i, c := range guild.Channels {
+			if c.ID == channel.ID {
+				guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...)
+				return nil
+			}
+		}
+	}
+
+	return errors.New("Channel not found.")
+}
+
+// GetGuildChannelById gets a channel by ID from a guild.
+func (s *State) GetGuildChannelByID(guildID string, channelID string) (*Channel, error) {
+	guild, err := s.GetGuildByID(guildID)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, c := range guild.Channels {
+		if c.ID == channelID {
+			return &c, nil
+		}
+	}
+	return nil, errors.New("Channel not found.")
+}
+
+// GetPrivateChannelByID gets a private channel by ID.
+func (s *State) GetPrivateChannelByID(channelID string) (*Channel, error) {
+	for _, c := range s.PrivateChannels {
+		if c.ID == channelID {
+			return &c, nil
+		}
+	}
+	return nil, errors.New("Channel not found.")
+}

+ 12 - 10
structs.go

@@ -84,6 +84,9 @@ type Session struct {
 	VChannelID string
 	Vop2       VoiceOP2
 	UDPConn    *net.UDPConn
+
+	// Managed state object, updated with events.
+	State *State
 }
 
 // A Message stores all data related to a specific Discord message.
@@ -262,14 +265,6 @@ type User struct {
 // it just doesn't seem able to handle this one
 // field correctly.  Need to research this more.
 
-// A PrivateChannel stores all data for a specific user private channel.
-type PrivateChannel struct {
-	ID            string `json:"id"`
-	IsPrivate     bool   `json:"is_private"`
-	LastMessageID string `json:"last_message_id"`
-	Recipient     User   `json:"recipient"`
-} // merge with channel?
-
 // A Settings stores data for a specific users Discord client settings.
 type Settings struct {
 	RenderEmbeds          bool     `json:"render_embeds"`
@@ -298,8 +293,8 @@ type Ready struct {
 	HeartbeatInterval time.Duration `json:"heartbeat_interval"`
 	User              User          `json:"user"`
 	ReadState         []ReadState
-	PrivateChannels   []PrivateChannel
-	Guilds            []Guild
+	PrivateChannels   []Channel `json:"private_channels"`
+	Guilds            []Guild   `json:"guilds"`
 }
 
 // A ReadState stores data on the read state of channels.
@@ -360,3 +355,10 @@ type GuildBan struct {
 	User    User   `json:"user"`
 	GuildID string `json:"guild_id"`
 }
+
+// A State contains the current known state, as discord sends this in a
+// READY blob, it seems reasonable to simply use that message type as the
+// data store.
+type State struct {
+	Ready
+}

+ 47 - 36
wsapi.go

@@ -153,6 +153,7 @@ func (s *Session) event(messageType int, message []byte) (err error) {
 	case "READY":
 		var st Ready
 		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.OnReady(&st)
 			if s.OnReady != nil {
 				s.OnReady(s, st)
 			}
@@ -235,77 +236,87 @@ func (s *Session) event(messageType int, message []byte) (err error) {
 			return
 		}
 	case "CHANNEL_CREATE":
-		if s.OnChannelCreate != nil {
-			var st Channel
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
+			fmt.Println("channel create", st)
+			s.State.AddChannel(&st)
+			if s.OnChannelCreate != nil {
 				s.OnChannelCreate(s, st)
 			}
-			return
 		}
+		return
 	case "CHANNEL_UPDATE":
-		if s.OnChannelUpdate != nil {
-			var st Channel
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.AddChannel(&st)
+			if s.OnChannelUpdate != nil {
 				s.OnChannelUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "CHANNEL_DELETE":
-		if s.OnChannelDelete != nil {
-			var st Channel
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.RemoveChannel(&st)
+			if s.OnChannelDelete != nil {
 				s.OnChannelDelete(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_CREATE":
-		if s.OnGuildCreate != nil {
-			var st Guild
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.AddGuild(&st)
+			if s.OnGuildCreate != nil {
 				s.OnGuildCreate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_UPDATE":
-		if s.OnGuildCreate != nil {
-			var st Guild
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.AddGuild(&st)
+			if s.OnGuildCreate != nil {
 				s.OnGuildUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_DELETE":
-		if s.OnGuildDelete != nil {
-			var st Guild
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.RemoveGuild(&st)
+			if s.OnGuildDelete != nil {
 				s.OnGuildDelete(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_ADD":
-		if s.OnGuildMemberAdd != nil {
-			var st Member
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.AddMember(&st)
+			if s.OnGuildMemberAdd != nil {
 				s.OnGuildMemberAdd(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_REMOVE":
-		if s.OnGuildMemberRemove != nil {
-			var st Member
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.RemoveMember(&st)
+			if s.OnGuildMemberRemove != nil {
 				s.OnGuildMemberRemove(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_UPDATE":
-		if s.OnGuildMemberUpdate != nil {
-			var st Member
-			if err = unmarshalEvent(e, &st); err == nil {
+		var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
+			s.State.AddMember(&st)
+			if s.OnGuildMemberUpdate != nil {
 				s.OnGuildMemberUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_ROLE_CREATE":
 		if s.OnGuildRoleCreate != nil {
 			var st GuildRole