add signature parser

multicall/signature.py

@@ -0,0 +1,39 @@
+from eth_abi import encode_single, decode_single
+from eth_utils import function_signature_to_4byte_selector
+
+
+def parse_signature(signature):
+    """
+    Breaks 'func(address)(uint256)' into ['func', '(address)', '(uint256)']
+    """
+    parts = []
+    stack = []
+    start = 0
+    for end, letter in enumerate(signature):
+        if letter == '(':
+            stack.append(letter)
+            if not parts:
+                parts.append(signature[start:end])
+                start = end
+        if letter == ')':
+            stack.pop()
+            if not stack:  # we are only interested in outermost groups
+                parts.append(signature[start:end + 1])
+                start = end + 1
+    return parts
+
+
+class Signature:
+    def __init__(self, signature):
+        self.signature = signature
+        self.parts = parse_signature(signature)
+        self.input_types = self.parts[1]
+        self.output_types = self.parts[2]
+        self.function = ''.join(self.parts[:2])
+        self.fourbyte = function_signature_to_4byte_selector(self.function)
+
+    def encode_data(self, args=None):
+        return self.fourbyte + encode_single(self.input_types, args) if args else self.fourbyte
+
+    def decode_data(self, output):
+        return decode_single(self.output_types, output)

tests/conftest.py

@@ -0,0 +1,4 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath('.'))

tests/test_signature.py

@@ -0,0 +1,23 @@
+from eth_abi import encode_abi, decode_abi
+from multicall import signature
+
+args = ((1, 2, 3), '0x' + 'f' * 40, b'data')
+types = ['uint256[]', 'address', 'bytes']
+
+
+def test_signature_parsing():
+    sig = signature.Signature('aggregate((address,bytes)[])(uint256,bytes[])')
+    assert sig.parts == ['aggregate', '((address,bytes)[])', '(uint256,bytes[])']
+    assert sig.input_types == '((address,bytes)[])'
+    assert sig.output_types == '(uint256,bytes[])'
+
+
+def test_signature_encoding():
+    sig = signature.Signature('test(uint256[],address,bytes)()')
+    assert sig.encode_data(args) == sig.fourbyte + encode_abi(types, args)
+
+
+def test_signature_decoding():
+    sig = signature.Signature('test()(uint256[],address,bytes)')
+    data = encode_abi(types, args)
+    assert sig.decode_data(data) == args