diff mercurial/wireprototypes.py @ 37781:352932a11905

wireproto: move command registration types to wireprototypes These are shared across wire protocol implementations. wireprototypes is our module for common code. Differential Revision: https://phab.mercurial-scm.org/D3396
author Gregory Szorc <gregory.szorc@gmail.com>
date Mon, 16 Apr 2018 21:52:33 -0700
parents 564a3eec6e63
children 9d818539abfa
line wrap: on
line diff
--- a/mercurial/wireprototypes.py	Mon Apr 16 21:49:59 2018 -0700
+++ b/mercurial/wireprototypes.py	Mon Apr 16 21:52:33 2018 -0700
@@ -241,3 +241,81 @@
         doesn't have that permission, the exception should raise or abort
         in a protocol specific manner.
         """
+
+class commandentry(object):
+    """Represents a declared wire protocol command."""
+    def __init__(self, func, args='', transports=None,
+                 permission='push'):
+        self.func = func
+        self.args = args
+        self.transports = transports or set()
+        self.permission = permission
+
+    def _merge(self, func, args):
+        """Merge this instance with an incoming 2-tuple.
+
+        This is called when a caller using the old 2-tuple API attempts
+        to replace an instance. The incoming values are merged with
+        data not captured by the 2-tuple and a new instance containing
+        the union of the two objects is returned.
+        """
+        return commandentry(func, args=args, transports=set(self.transports),
+                            permission=self.permission)
+
+    # Old code treats instances as 2-tuples. So expose that interface.
+    def __iter__(self):
+        yield self.func
+        yield self.args
+
+    def __getitem__(self, i):
+        if i == 0:
+            return self.func
+        elif i == 1:
+            return self.args
+        else:
+            raise IndexError('can only access elements 0 and 1')
+
+class commanddict(dict):
+    """Container for registered wire protocol commands.
+
+    It behaves like a dict. But __setitem__ is overwritten to allow silent
+    coercion of values from 2-tuples for API compatibility.
+    """
+    def __setitem__(self, k, v):
+        if isinstance(v, commandentry):
+            pass
+        # Cast 2-tuples to commandentry instances.
+        elif isinstance(v, tuple):
+            if len(v) != 2:
+                raise ValueError('command tuples must have exactly 2 elements')
+
+            # It is common for extensions to wrap wire protocol commands via
+            # e.g. ``wireproto.commands[x] = (newfn, args)``. Because callers
+            # doing this aren't aware of the new API that uses objects to store
+            # command entries, we automatically merge old state with new.
+            if k in self:
+                v = self[k]._merge(v[0], v[1])
+            else:
+                # Use default values from @wireprotocommand.
+                v = commandentry(v[0], args=v[1],
+                                 transports=set(TRANSPORTS),
+                                 permission='push')
+        else:
+            raise ValueError('command entries must be commandentry instances '
+                             'or 2-tuples')
+
+        return super(commanddict, self).__setitem__(k, v)
+
+    def commandavailable(self, command, proto):
+        """Determine if a command is available for the requested protocol."""
+        assert proto.name in TRANSPORTS
+
+        entry = self.get(command)
+
+        if not entry:
+            return False
+
+        if proto.name not in entry.transports:
+            return False
+
+        return True