emmental/lib/tagstore.py

# Copyright 2021 (c) Anna Schumaker.
from . import publisher
from . import tag
import threading

class TagStore:
    def __init__(self):
        self.store = dict()
        self.lock = threading.Lock()
        self.Added = publisher.Publisher()
        self.Removed = publisher.Publisher()

    def __getitem__(self, name):
        with self.lock:
            return self.store.get(name)

    def __getstate__(self):
        with self.lock:
            return { "store" : self.store }

    def __get_tag__(self, name):
        with self.lock:
            if (t := self.store.get(name)) != None:
                return t
            t = tag.Tag(name)
            self.store[name] = t
        self.Added.publish(t)
        return t

    def __len__(self):
        with self.lock:
            return len(self.store)

    def __setstate__(self, state):
        self.store = state["store"]
        self.lock = threading.Lock()
        self.Added = publisher.Publisher()
        self.Removed = publisher.Publisher()

    def __setstate__(self, state):
        self.store = state["store"]
        self.lock = threading.Lock()
        self.Added = publisher.Publisher()
        self.Removed = publisher.Publisher()

    def add(self, name, track=None):
        t = self.__get_tag__(name.strip())
        if track:
            t.add_track(track)
        return t

    def init_track(self, name, track):
        with self.lock:
            if (t := self.store.get(name)) != None:
                t.init_track(track)
                return t

    def remove(self, t, track=None):
        if track:
            t.remove_track(track)
        if (track == None or len(t) == 0) and t in self.store.values():
            with self.lock:
                self.store.pop(t.name)
            self.Removed.publish(t)

    def reset(self):
        with self.lock:
            self.store.clear()
            self.Added.reset()
            self.Removed.reset()

    def tags(self):
        with self.lock:
            for (name, tag) in self.store.items():
                yield tag