anidb: Add a few more linking methods

[go-anidb.git] / misc / episodelist.go

diff --git a/misc/episodelist.go b/misc/episodelist.go
index a22bef488b942af3fc17003914185b660c500bce..a0fb2a0660efd018ee2166d3496fc00864dfe7a9 100644 (file)
--- a/misc/episodelist.go
+++ b/misc/episodelist.go
@@ -1,6 +1,7 @@
 package misc
 
 import (
+       "encoding/json"
        "fmt"
        "sort"
        "strings"
@@ -17,6 +18,27 @@ type EpisodeCount struct {
 
 type EpisodeList []*EpisodeRange
 
+func EpisodeToList(ep *Episode) EpisodeList {
+       return RangesToList(EpisodeToRange(ep))
+}
+
+func RangesToList(ranges ...*EpisodeRange) EpisodeList {
+       return EpisodeList(ranges)
+}
+
+func ContainerToList(ec EpisodeContainer) EpisodeList {
+       switch v := ec.(type) {
+       case *Episode:
+               return EpisodeToList(v)
+       case *EpisodeRange:
+               return RangesToList(v)
+       case EpisodeList:
+               return v
+       default:
+               panic("unimplemented")
+       }
+}
+
 // Converts the EpisodeList into the AniDB API list format.
 func (el EpisodeList) String() string {
        scales := map[EpisodeType]int{}
@@ -66,6 +88,46 @@ func (el EpisodeList) FormatLog(ec EpisodeCount) string {
        return strings.Join(parts, ",")
 }
 
+func (el EpisodeList) Infinite() bool {
+       for i := range el {
+               if el[i].Infinite() {
+                       return true
+               }
+       }
+       return false
+}
+
+// Returns a channel that can be used to iterate using for/range.
+//
+// If the EpisodeList is infinite, then the channel is also infinite.
+// The caller is allowed to close the channel in such case.
+//
+// NOTE: Not thread safe.
+func (el EpisodeList) Episodes() chan Episode {
+       ch := make(chan Episode, 1)
+
+       go func() {
+               abort := false
+
+               if el.Infinite() {
+                       defer func() { recover(); abort = true }()
+               } else {
+                       defer close(ch)
+               }
+
+               for _, er := range el {
+                       for ep := range er.Episodes() {
+                               ch <- ep
+
+                               if abort {
+                                       return
+                               }
+                       }
+               }
+       }()
+       return ch
+}
+
 // Returns true if any of the contained EpisodeRanges contain the
 // given EpisodeContainer.
 func (el EpisodeList) ContainsEpisodes(ec EpisodeContainer) bool {
@@ -134,6 +196,37 @@ func (el EpisodeList) Simplify() EpisodeList {
        return nl
 }
 
+func (el EpisodeList) CountEpisodes() (ec EpisodeCount) {
+       for _, er := range el {
+               var c *int
+               switch er.Type {
+               case EpisodeTypeRegular:
+                       c = &ec.RegularCount
+               case EpisodeTypeSpecial:
+                       c = &ec.SpecialCount
+               case EpisodeTypeCredits:
+                       c = &ec.CreditsCount
+               case EpisodeTypeOther:
+                       c = &ec.OtherCount
+               case EpisodeTypeTrailer:
+                       c = &ec.TrailerCount
+               case EpisodeTypeParody:
+                       c = &ec.ParodyCount
+               default:
+                       continue
+               }
+               if *c < 0 {
+                       continue
+               }
+               if er.End == nil {
+                       *c = -1
+                       continue
+               }
+               *c += er.End.Number - er.Start.Number
+       }
+       return
+}
+
 func (el EpisodeList) Len() int {
        return len(el)
 }
@@ -157,3 +250,55 @@ func (el EpisodeList) Less(i, j int) bool {
 func (el EpisodeList) Swap(i, j int) {
        el[i], el[j] = el[j], el[i]
 }
+
+func (el *EpisodeList) Add(ec EpisodeContainer) {
+       *el = append(*el, ContainerToList(ec)...)
+       *el = el.Simplify()
+}
+
+func (el *EpisodeList) Sub(ec EpisodeContainer) {
+       el2 := make(EpisodeList, 0, len(*el)*2)
+       switch e, ok := ec.(canInfinite); {
+       case ok:
+               if e.Infinite() {
+                       eCh := e.Episodes()
+                       ep := <-eCh
+                       close(eCh)
+
+                       for _, r := range *el {
+                               el2 = append(el2, r.Split(&ep)[0])
+                       }
+                       break
+               }
+               fallthrough
+       default:
+               for ep := range ec.Episodes() {
+                       for _, r := range *el {
+                               el2 = append(el2, r.Split(&ep)...)
+                       }
+                       el2 = el2.Simplify()
+               }
+       }
+       *el = append(*el, el2.Simplify()...)
+}
+
+// Equivalent to marshaling el.String()
+func (el EpisodeList) MarshalJSON() ([]byte, error) {
+       return json.Marshal(el.String())
+}
+
+// NOTE: Since the String() representation doesn't include them,
+// it's not exactly reversible if the user has set .Parts in any
+// of the contained episodes.
+func (el EpisodeList) UnmarshalJSON(b []byte) error {
+       var v string
+       if err := json.Unmarshal(b, &v); err != nil {
+               return err
+       }
+
+       l := ParseEpisodeList(v)
+       for k := range l {
+               el[k] = l[k]
+       }
+       return nil
+}