Merge pull request #42 from iopred/state

Implement state tracking.

discord.go

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

state.go

@@ -0,0 +1,268 @@
+package discordgo
+
+import "errors"
+
+var nilError error = errors.New("State not instantiated, please use discordgo.New() or assign session.State.")
+
+// 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) error {
+	if s == nil {
+		return nilError
+	}
+
+	s.Ready = *r
+	return nil
+}
+
+// GuildAdd adds a guild to the current world state, or
+// updates it if it already exists.
+func (s *State) GuildAdd(guild *Guild) error {
+	if s == nil {
+		return nilError
+	}
+
+	for _, g := range s.Guilds {
+		if g.ID == guild.ID {
+			// This could be a little faster ;)
+			for _, m := range guild.Members {
+				s.MemberAdd(&m)
+			}
+			for _, c := range guild.Channels {
+				s.ChannelAdd(&c)
+			}
+			return nil
+		}
+	}
+	s.Guilds = append(s.Guilds, *guild)
+	return nil
+}
+
+// GuildRemove removes a guild from current world state.
+func (s *State) GuildRemove(guild *Guild) error {
+	if s == nil {
+		return nilError
+	}
+
+	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.")
+}
+
+// Guild gets a guild by ID.
+// Useful for querying if @me is in a guild:
+//     _, err := discordgo.Session.State.Guild(guildID)
+//     isInGuild := err == nil
+func (s *State) Guild(guildID string) (*Guild, error) {
+	if s == nil {
+		return nil, nilError
+	}
+
+	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.
+
+// MemberAdd adds a member to the current world state, or
+// updates it if it already exists.
+func (s *State) MemberAdd(member *Member) error {
+	if s == nil {
+		return nilError
+	}
+
+	guild, err := s.Guild(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
+}
+
+// MemberRemove removes a member from current world state.
+func (s *State) MemberRemove(member *Member) error {
+	if s == nil {
+		return nilError
+	}
+
+	guild, err := s.Guild(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.")
+}
+
+// Member gets a member by ID from a guild.
+func (s *State) Member(guildID string, userID string) (*Member, error) {
+	if s == nil {
+		return nil, nilError
+	}
+
+	guild, err := s.Guild(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.")
+}
+
+// ChannelAdd 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) ChannelAdd(channel *Channel) error {
+	if s == nil {
+		return nilError
+	}
+
+	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.Guild(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
+}
+
+// ChannelRemove removes a channel from current world state.
+func (s *State) ChannelRemove(channel *Channel) error {
+	if s == nil {
+		return nilError
+	}
+
+	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.Guild(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.")
+}
+
+// GuildChannel gets a channel by ID from a guild.
+func (s *State) GuildChannel(guildID string, channelID string) (*Channel, error) {
+	if s == nil {
+		return nil, nilError
+	}
+
+	guild, err := s.Guild(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.")
+}
+
+// PrivateChannel gets a private channel by ID.
+func (s *State) PrivateChannel(channelID string) (*Channel, error) {
+	if s == nil {
+		return nil, nilError
+	}
+
+	for _, c := range s.PrivateChannels {
+		if c.ID == channelID {
+			return &c, nil
+		}
+	}
+
+	return nil, errors.New("Channel not found.")
+}
+
+// Channel gets a channel by ID, it will look in all guilds an private channels.
+func (s *State) Channel(channelID string) (*Channel, error) {
+	if s == nil {
+		return nil, nilError
+	}
+
+	c, err := s.PrivateChannel(channelID)
+	if err == nil {
+		return c, nil
+	}
+
+	for _, g := range s.Guilds {
+		c, err := s.GuildChannel(g.ID, channelID)
+		if err == nil {
+			return c, nil
+		}
+	}
+
+	return nil, errors.New("Channel not found.")
+}

structs.go

@@ -84,6 +84,10 @@ type Session struct {
 	VChannelID string
 	Vop2       VoiceOP2
 	UDPConn    *net.UDPConn
+
+	// Managed state object, updated with events.
+	State        *State
+	StateEnabled bool
 }
 
 // A Message stores all data related to a specific Discord message.
@@ -262,14 +266,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 +294,8 @@ type Ready struct {
 	HeartbeatInterval time.Duration `json:"heartbeat_interval"`
 	User              User          `json:"user"`
 	ReadState         []ReadState
-	PrivateChannels   []PrivateChannel
+	PrivateChannels   []Channel `json:"private_channels"`
-	Guilds            []Guild
+	Guilds            []Guild   `json:"guilds"`
 }
 
 // A ReadState stores data on the read state of channels.
@@ -360,3 +356,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 struct as the data store.
+type State struct {
+	Ready
+}

wsapi.go

@@ -153,6 +153,9 @@ func (s *Session) event(messageType int, message []byte) (err error) {
 	case "READY":
 		var st Ready
 		if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.OnReady(&st)
+			}
 			if s.OnReady != nil {
 				s.OnReady(s, st)
 			}
@@ -235,77 +238,104 @@ func (s *Session) event(messageType int, message []byte) (err error) {
 			return
 		}
 	case "CHANNEL_CREATE":
-		if s.OnChannelCreate != nil {
+		var st Channel
-			var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.ChannelAdd(&st)
+			}
+			if s.OnChannelCreate != nil {
 				s.OnChannelCreate(s, st)
 			}
-			return
 		}
+		return
 	case "CHANNEL_UPDATE":
-		if s.OnChannelUpdate != nil {
+		var st Channel
-			var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.ChannelAdd(&st)
+			}
+			if s.OnChannelUpdate != nil {
 				s.OnChannelUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "CHANNEL_DELETE":
-		if s.OnChannelDelete != nil {
+		var st Channel
-			var st Channel
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.ChannelRemove(&st)
+			}
+			if s.OnChannelDelete != nil {
 				s.OnChannelDelete(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_CREATE":
-		if s.OnGuildCreate != nil {
+		var st Guild
-			var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.GuildAdd(&st)
+			}
+			if s.OnGuildCreate != nil {
 				s.OnGuildCreate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_UPDATE":
-		if s.OnGuildCreate != nil {
+		var st Guild
-			var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.GuildAdd(&st)
+			}
+			if s.OnGuildCreate != nil {
 				s.OnGuildUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_DELETE":
-		if s.OnGuildDelete != nil {
+		var st Guild
-			var st Guild
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.GuildRemove(&st)
+			}
+			if s.OnGuildDelete != nil {
 				s.OnGuildDelete(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_ADD":
-		if s.OnGuildMemberAdd != nil {
+		var st Member
-			var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.MemberAdd(&st)
+			}
+			if s.OnGuildMemberAdd != nil {
 				s.OnGuildMemberAdd(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_REMOVE":
-		if s.OnGuildMemberRemove != nil {
+		var st Member
-			var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.MemberRemove(&st)
+			}
+			if s.OnGuildMemberRemove != nil {
 				s.OnGuildMemberRemove(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_MEMBER_UPDATE":
-		if s.OnGuildMemberUpdate != nil {
+		var st Member
-			var st Member
+		if err = unmarshalEvent(e, &st); err == nil {
-			if err = unmarshalEvent(e, &st); err == nil {
+			if s.StateEnabled {
+				s.State.MemberAdd(&st)
+			}
+			if s.OnGuildMemberUpdate != nil {
 				s.OnGuildMemberUpdate(s, st)
 			}
-			return
 		}
+		return
 	case "GUILD_ROLE_CREATE":
 		if s.OnGuildRoleCreate != nil {
 			var st GuildRole