Added Tree types and get_bugdir_subscribers to becommands/subscribe.py.

The Tree subclass SubscriptionType may be a few more lines to setup,
but they should be really easy to extend and will automatically keep
the longhelp and type handling in sync.

An example extension for bugdir types would be
  all -> assigned -> <user>
to subscribe only to bugs being assigned to the specified user.  You'd
have to loosen the currently strict InvalidType checking to make that
work, but the current type-tree handling would be up to the task.

Also a bit of reorganizing to hide the private functions.

1 files changed, 214 insertions, 83 deletions

diff --git a/becommands/subscribe.py b/becommands/subscribe.py
index 26158b1..64a2867 100644
--- a/becommands/subscribe.py
+++ b/becommands/subscribe.py
@@ -14,15 +14,49 @@
 # with this program; if not, write to the Free Software Foundation, Inc.,
 # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 """(Un)subscribe to change notification"""
-from libbe import cmdutil, bugdir
+from libbe import cmdutil, bugdir, tree
 import os, copy
 __desc__ = __doc__
 
 TAG="SUBSCRIBE:"
 
+class SubscriptionType (tree.Tree):
+    """
+    Trees of subscription types to allow users to select exactly what
+    notifications they want to subscribe to.
+    """
+    def __init__(self, type_name, *args, **kwargs):
+        tree.Tree.__init__(self, *args, **kwargs)
+        self.type = type_name
+    def __str__(self):
+        return self.type
+    def __repr__(self):
+        return "<SubscriptionType: %s>" % str(self)
+    def string_tree(self, indent=0):
+        lines = []
+        for depth,node in self.thread():
+            lines.append("%s%s" % (" "*(indent+2*depth), node))
+        return "\n".join(lines)
+
+BUGDIR_TYPE_NEW = SubscriptionType("new")
+BUGDIR_TYPE_ALL = SubscriptionType("all", [BUGDIR_TYPE_NEW])
+
+# same name as BUGDIR_TYPE_ALL for consistency
+BUG_TYPE_ALL = SubscriptionType(str(BUGDIR_TYPE_ALL))
+
+INVALID_TYPE = SubscriptionType("INVALID")
+
+class InvalidType (ValueError):
+    def __init__(self, type_name, type_root):
+        msg = "Invalid type %s for tree:\n%s" \
+            % (type_name, type_root.string_tree(4))
+        ValueError.__init__(self, msg)
+        self.type_name = type_name
+        self.type_root = type_root
+
+
 def execute(args, manipulate_encodings=True):
     """
-    >>> from libbe import utility
     >>> bd = bugdir.simple_bug_dir()
     >>> bd.set_sync_with_disk(True)
     >>> os.chdir(bd.root)
@@ -92,23 +126,23 @@ def execute(args, manipulate_encodings=True):
     types = options.types.split(",")
 
     if len(args) == 0 or args[0] == "DIR": # directory-wide subscriptions
-        if options.unsubscribe == False:
-            for t in types:
-                assert t in ["all","new"], t
+        type_root = BUGDIR_TYPE_ALL
         entity = bd
         entity_name = "bug directory"
     else: # bug-specific subscriptions
-        if options.unsubscribe == False:
-            assert types == ["all"], types
+        type_root = BUG_TYPE_ALL
         bug = bd.bug_from_shortname(args[0])
         entity = bug
         entity_name = bug.uuid
 
+    types = [type_from_name(name, type_root, default=INVALID_TYPE,
+                            default_ok=options.unsubscribe)
+             for name in types]
     estrs = entity.extra_strings
     if options.unsubscribe == True:
-        estrs = unsubscribe(estrs, subscriber, types, servers)
+        estrs = unsubscribe(estrs, subscriber, types, servers, type_root)
     else: # add the tag
-        estrs = subscribe(estrs, subscriber, types, servers)
+        estrs = subscribe(estrs, subscriber, types, servers, type_root)
     entity.extra_strings = estrs # reassign to notice change
 
     subscriptions = []
@@ -120,67 +154,120 @@ def execute(args, manipulate_encodings=True):
         print "Subscriptions for %s:" % entity_name
         print '\n'.join(subscriptions)
 
-def generate_string(subscriber, types, servers):
+
+def get_parser():
+    parser = cmdutil.CmdOptionParser("be subscribe ID")
+    parser.add_option("-u", "--unsubscribe", action="store_true",
+                      dest="unsubscribe", default=False,
+                      help="Unsubscribe instead of subscribing.")
+    parser.add_option("-s", "--subscriber", dest="subscriber",
+                      metavar="SUBSCRIBER",
+                      help="Email address of the subscriber (defaults to bugdir.user_id).")
+    parser.add_option("-S", "--servers", dest="servers", metavar="SERVERS",
+                      help="Servers from which you want notification.")
+    parser.add_option("-t", "--type", dest="types", metavar="TYPES",
+                      help="Types of changes you wish to be notified about.")
+    return parser
+
+longhelp="""
+ID can be either a bug id, or blank/"DIR", in which case it refers to the
+whole bug directory.
+
+SERVERS specifies the servers from which you would like to receive
+notification.  Multiple severs may be specified in a comma-separated
+list, or you can use "*" to match all servers (the default).  If you
+have not selected a server, it should politely refrain from notifying
+you of changes, although there is no way to guarantee this behavior.
+
+Available TYPES:
+  For bugs:
+%s
+  For DIR :
+%s
+
+For unsubscription, any listed SERVERS and TYPES are removed from your
+subscription.  Either the catch-all server "*" or type "%s" will
+remove SUBSCRIBER entirely from the specified ID.
+
+This command is intended for use primarily by public interfaces, since
+if you're just hacking away on your private repository, you'll known
+what's changed ;).  This command just (un)sets the appropriate
+subscriptions, and leaves it up to each interface to perform the
+notification.
+""" % (BUG_TYPE_ALL.string_tree(6), BUGDIR_TYPE_ALL.string_tree(6),
+       BUGDIR_TYPE_ALL)
+
+def help():
+    return get_parser().help_str() + longhelp
+
+# internal helper functions
+
+def _generate_string(subscriber, types, servers):
+    types = sorted([str(t) for t in types])
+    servers = sorted(servers)
     return "%s%s\t%s\t%s" % (TAG,subscriber,",".join(types),",".join(servers))
 
-def parse_string(string):
+def _parse_string(string, type_root):
     assert string.startswith(TAG), string
     string = string[len(TAG):]
     subscriber,types,servers = string.split("\t")
-    return (subscriber,types.split(","),servers.split(","))
+    types = [type_from_name(name, type_root) for name in types.split(",")]
+    return (subscriber,types,servers.split(","))
 
-def get_subscribers(extra_strings, type, server):
-    for string in extra_strings:
-        subscriber,types,servers = parse_string(string)
-        type_match = False
-        if type in types or types == ["all"]:
-            type_match = True
-        server_match = False
-        if server in servers or servers == ["*"]:
-            server_match = True
-        if type_match == True and server_match == True: 
-            yield subscriber
-
-def get_matching_string(extra_strings, subscriber, types, servers):
+def _get_subscriber(extra_strings, subscriber, type_root):
     for i,string in enumerate(extra_strings):
         if string.startswith(TAG):
-            s,ts,srvs = parse_string(string)
+            s,ts,srvs = _parse_string(string, type_root)
             if s == subscriber:
                 return i,s,ts,srvs # match!
     return None # no match
 
-def subscribe(extra_strings, subscriber, types, servers):
-    args = get_matching_string(extra_strings, subscriber, types, servers)
+# functions exposed to other modules
+
+def type_from_name(name, type_root, default=None, default_ok=False):
+    if name == str(type_root):
+        return type_root
+    for t in type_root.traverse():
+        if name == str(t):
+            return t
+    if default_ok:
+        return default
+    raise InvalidType(name, type_root)
+
+def subscribe(extra_strings, subscriber, types, servers, type_root):
+    args = _get_subscriber(extra_strings, subscriber, type_root)
     if args == None: # no match
-        extra_strings.append(generate_string(subscriber, types, servers))
+        extra_strings.append(_generate_string(subscriber, types, servers))
         return extra_strings
     # Alter matched string
     i,s,ts,srvs = args
-    if "all" in types+ts:
-        ts = ["all"]
-    else:
-        ts = list(set(types+ts))
-        ts.sort()
+    for t in types:
+        if t not in ts:
+            ts.append(t)
+    # remove descendant types
+    all_ts = copy.copy(ts)
+    for t in all_ts:
+        for tt in all_ts:
+            if tt in ts and t.has_descendant(tt):
+                ts.remove(tt)
     if "*" in servers+srvs:
         srvs = ["*"]
     else:
         srvs = list(set(servers+srvs))
-        srvs.sort()
-    extra_strings[i] = generate_string(subscriber, ts, srvs)
+    extra_strings[i] = _generate_string(subscriber, ts, srvs)
     return extra_strings
 
-def unsubscribe(extra_strings, subscriber, types, servers):
-    args = get_matching_string(extra_strings, subscriber, types, servers)
+def unsubscribe(extra_strings, subscriber, types, servers, type_root):
+    args = _get_subscriber(extra_strings, subscriber, type_root)
     if args == None: # no match
         return extra_strings # pass
     # Remove matched string
     i,s,ts,srvs = args
-    if "all" in types:
-        ts = []
-    else:
-        for t in types:
-            if t in ts:
-                ts.remove(t)
+    all_ts = copy.copy(ts)
+    for t in types:
+        for tt in all_ts:
+            if tt in ts and t.has_descendant(tt):
+                ts.remove(tt)
     if "*" in servers+srvs:
         srvs = []
     else:
@@ -190,48 +277,92 @@ def unsubscribe(extra_strings, subscriber, types, servers):
     if len(ts) == 0 or len(srvs) == 0:
         extra_strings.pop(i)
     else:
-        extra_strings[i] = generate_string(subscriber, ts, srvs)
+        extra_strings[i] = _generate_string(subscriber, ts, srvs)
     return extra_strings
 
-def get_parser():
-    parser = cmdutil.CmdOptionParser("be subscribe ID")
-    parser.add_option("-u", "--unsubscribe", action="store_true",
-                      dest="unsubscribe", default=False,
-                      help="Unsubscribe instead of subscribing.")
-    parser.add_option("-s", "--subscriber", dest="subscriber",
-                      metavar="SUBSCRIBER",
-                      help="Email address of the subscriber (defaults to bugdir.user_id).")
-    parser.add_option("-S", "--servers", dest="servers", metavar="SERVERS",
-                      help="Servers from which you want notification.")
-    parser.add_option("-t", "--type", dest="types", metavar="TYPES",
-                      help="Types of changes you wish to be notified about.")
-    return parser
-
-longhelp="""
-ID can be either a bug id, or blank/"DIR", in which case it refers to the
-whole bug directory.
-
-SERVERS specifies the servers from which you would like to receive
-notification.  Multiple severs may be specified in a comma-separated
-list, or you can use "*" to match all servers (the default).  If you
-have not selected a server, it should politely refrain from notifying
-you of changes, although there is no way to guarantee this behavior.
+def get_subscribers(extra_strings, type, server, type_root,
+                    match_ancestor_types=False,
+                    match_descendant_types=False):
+    """
+    Set match_ancestor_types=True if you want to find eveyone who
+    cares about your particular type.
 
-Available TYPES:
-  For bugs:  all
-  For DIR :  all
-             new  - only notify when new bugs are added
+    Set match_descendant_types=True if you want to find subscribers
+    who may only care about some subset of your type.  This is useful
+    for generating lists of all the subscribers in a given set of
+    extra_strings.
 
-For unsubscription, any listed SERVERS and TYPES are removed from your
-subscription.  Either the catch-all server "*" or type "all" will
-remove SUBSCRIBER entirely from the specified ID.
+    >>> def sgs(*args, **kwargs):
+    ...     return sorted(get_subscribers(*args, **kwargs))
+    >>> es = []
+    >>> es = subscribe(es, "John Doe <j@doe.com>", [BUGDIR_TYPE_ALL], ["a.com"], BUGDIR_TYPE_ALL)
+    >>> es = subscribe(es, "Jane Doe <J@doe.com>", [BUGDIR_TYPE_NEW], ["*"], BUGDIR_TYPE_ALL)
+    >>> sgs(es, BUGDIR_TYPE_ALL, "a.com", BUGDIR_TYPE_ALL)
+    ['John Doe <j@doe.com>']
+    >>> sgs(es, BUGDIR_TYPE_ALL, "a.com", BUGDIR_TYPE_ALL, match_descendant_types=True)
+    ['Jane Doe <J@doe.com>', 'John Doe <j@doe.com>']
+    >>> sgs(es, BUGDIR_TYPE_ALL, "b.net", BUGDIR_TYPE_ALL, match_descendant_types=True)
+    ['Jane Doe <J@doe.com>']
+    >>> sgs(es, BUGDIR_TYPE_NEW, "a.com", BUGDIR_TYPE_ALL)
+    ['Jane Doe <J@doe.com>']
+    >>> sgs(es, BUGDIR_TYPE_NEW, "a.com", BUGDIR_TYPE_ALL, match_ancestor_types=True)
+    ['Jane Doe <J@doe.com>', 'John Doe <j@doe.com>']
+    """
+    for string in extra_strings:
+        subscriber,types,servers = _parse_string(string, type_root)
+        type_match = False
+        if type in types:
+            type_match = True
+        if type_match == False and match_ancestor_types == True:
+            for t in types:
+                if t.has_descendant(type):
+                    type_match = True
+                    break
+        if type_match == False and match_descendant_types == True:
+            for t in types:
+                if type.has_descendant(t):
+                    type_match = True
+                    break
+        server_match = False
+        if server in servers or servers == ["*"]:
+            server_match = True
+        if type_match == True and server_match == True:
+            yield subscriber
 
-This command is intended for use primarily by public interfaces, since
-if you're just hacking away on your private repository, you'll known
-what's changed ;).  This command just (un)sets the appropriate
-subscriptions, and leaves it up to each interface to perform the
-notification.
-"""
+def get_bugdir_subscribers(bugdir, server):
+    """
+    I have a bugdir.  Who cares about it, and what do they care about?
+    Returns a dict of dicts:
+      subscribers[user][id] = types
+    where id is either a bug.uuid (in the case of a bug subscription)
+    or "DIR" (in the case of a bugdir subscription).
 
-def help():
-    return get_parser().help_str() + longhelp
+    >>> bd = bugdir.simple_bug_dir()
+    >>> a = bd.bug_from_shortname("a")
+    >>> bd.extra_strings = subscribe(bd.extra_strings, "John Doe <j@doe.com>", [BUGDIR_TYPE_ALL], ["a.com"], BUGDIR_TYPE_ALL)
+    >>> bd.extra_strings = subscribe(bd.extra_strings, "Jane Doe <J@doe.com>", [BUGDIR_TYPE_NEW], ["*"], BUGDIR_TYPE_ALL)
+    >>> a.extra_strings = subscribe(a.extra_strings, "John Doe <j@doe.com>", [BUG_TYPE_ALL], ["a.com"], BUG_TYPE_ALL)
+    >>> subscribers = get_bugdir_subscribers(bd, "a.com")
+    >>> subscribers["Jane Doe <J@doe.com>"]["DIR"]
+    [<SubscriptionType: new>]
+    >>> subscribers["John Doe <j@doe.com>"]["DIR"]
+    [<SubscriptionType: all>]
+    >>> subscribers["John Doe <j@doe.com>"]["a"]
+    [<SubscriptionType: all>]
+    >>> get_bugdir_subscribers(bd, "b.net")
+    {'Jane Doe <J@doe.com>': {'DIR': [<SubscriptionType: new>]}}
+    """
+    subscribers = {}
+    for sub in get_subscribers(bugdir.extra_strings, BUGDIR_TYPE_ALL, server,
+                               BUGDIR_TYPE_ALL, match_descendant_types=True):
+        i,s,ts,srvs = _get_subscriber(bugdir.extra_strings,sub,BUGDIR_TYPE_ALL)
+        subscribers[sub] = {"DIR":ts}
+    for bug in bugdir:
+        for sub in get_subscribers(bug.extra_strings, BUG_TYPE_ALL, server,
+                                   BUG_TYPE_ALL, match_descendant_types=True):
+            i,s,ts,srvs = _get_subscriber(bug.extra_strings,sub,BUG_TYPE_ALL)
+            if sub in subscribers:
+                subscribers[sub][bug.uuid] = ts
+            else:
+                subscribers[sub] = {bug.uuid:ts}
+    return subscribers