Map literals and tests

WDL/Expr.py

@@ -359,7 +359,6 @@ def eval(self, env : Env.Values) -> V.Base:
             raise Error.UnknownIdentifier(self) from None
 
 # Pair literal
-
 class Pair(Base):
     left : Base
     right : Base
@@ -379,3 +378,38 @@ def eval(self, env : Env.Values) -> V.Base:
         lv = self.left.eval(env)
         rv = self.right.eval(env)
         return V.Pair(self.type, (lv,rv))
+
+# Map literal
+class Map(Base):
+    items : List[Tuple[Base,Base]]
+
+    def __init__(self, pos : SourcePosition, items : List[Tuple[Base,Base]]) -> None:
+        super().__init__(pos)
+        self.items = items
+
+    def _infer_type(self, type_env : Env.Types) -> T.Base:
+        kty = None
+        vty = None
+        for k,v in self.items:
+            k.infer_type(type_env)
+            if kty is None:
+                kty = k.type
+            else:
+                k.typecheck(kty)
+            v.infer_type(type_env)
+            if vty is None or vty == T.Array(None) or vty == T.Map(None):
+                vty = v.type
+            else:
+                v.typecheck(vty)
+        if kty is None:
+            return T.Map(None)
+        assert vty is not None
+        return T.Map((kty,vty))
+
+    def eval(self, env : Env.Values) -> V.Base:
+        assert isinstance(self.type, T.Map)
+        eitems = []
+        for k,v in self.items:
+            eitems.append((k.eval(env), v.eval(env)))
+        # TODO: complain of duplicate keys
+        return V.Map(self.type, eitems)

WDL/StdLib.py

@@ -8,30 +8,62 @@
 import copy
 
 # Special function for array access arr[index], returning the element type
-class _ArrayGet(E._Function):
+#                      or map access map[key], returning the value type
+class _Get(E._Function):
     def infer_type(self, expr : E.Apply) -> T.Base:
         assert len(expr.arguments) == 2
-        if not isinstance(expr.arguments[0].type, T.Array):
-            raise Error.NotAnArray(expr.arguments[0])
-        if expr.arguments[0].type.item_type is None:
-            # the user wrote: [][idx]
-            raise Error.OutOfBounds(expr)
-        try:
-            expr.arguments[1].typecheck(T.Int())
-        except Error.StaticTypeMismatch:
-            raise Error.StaticTypeMismatch(expr.arguments[1], T.Int(), expr.arguments[1].type, "Array index") from None
-        return expr.arguments[0].type.item_type
+        lhs = expr.arguments[0]
+        rhs = expr.arguments[1]
+        if isinstance(lhs.type, T.Array):
+            if lhs.type.item_type is None:
+                # the user wrote: [][idx]
+                raise Error.OutOfBounds(expr)
+            try:
+                rhs.typecheck(T.Int())
+            except Error.StaticTypeMismatch:
+                raise Error.StaticTypeMismatch(rhs, T.Int(), rhs.type, "Array index") from None
+            return lhs.type.item_type
+        elif isinstance(lhs.type, T.Map):
+            if lhs.type.item_type is None:
+                raise Error.OutOfBounds(expr)
+            try:
+                rhs.typecheck(lhs.type.item_type[0])
+            except Error.StaticTypeMismatch:
+                raise Error.StaticTypeMismatch(rhs, lhs.type.item_type[0], rhs.type, "Map key") from None
+            return lhs.type.item_type[1]
+        else:
+            raise Error.NotAnArray(lhs)
 
     def __call__(self, expr : E.Apply, env : E.Env) -> V.Base:
         assert len(expr.arguments) == 2
-        arr = expr.arguments[0].eval(env)
-        assert isinstance(arr.type, T.Array)
-        assert isinstance(arr.value, list)
-        idx = expr.arguments[1].eval(env).expect(T.Int()).value
-        if idx < 0 or idx >= len(arr.value):
-            raise Error.OutOfBounds(expr.arguments[1])
-        return arr.value[idx] # pyre-ignore
-E._stdlib["_get"] = _ArrayGet()
+        lhs = expr.arguments[0]
+        rhs = expr.arguments[1]
+        if isinstance(lhs.type, T.Array): # pyre-fixme
+            arr = lhs.eval(env)
+            assert isinstance(arr, V.Array)
+            assert isinstance(arr.type, T.Array)
+            assert isinstance(arr.value, list)
+            idx = rhs.eval(env).expect(T.Int()).value
+            if idx < 0 or idx >= len(arr.value):
+                raise Error.OutOfBounds(rhs)
+            return arr.value[idx] # pyre-fixme
+        elif isinstance(lhs.type, T.Map):
+            mp = lhs.eval(env)
+            assert isinstance(mp, V.Map)
+            assert isinstance(mp.type, T.Map)
+            assert mp.type.item_type is not None
+            assert isinstance(mp.value, list)
+            ans = None
+            key = rhs.eval(env).expect(mp.type.item_type[0])
+            for k,v in mp.value:
+                if key == k:
+                    ans = v.expect(mp.type.item_type[1])
+            if ans is None:
+                raise Error.OutOfBounds(rhs) # TODO: KeyNotFound
+            return ans # pyre-fixme
+        else:
+            assert False
+E._stdlib["_get"] = _Get()
 
 # Pair get (EXPR.left/EXPR.right)
 # The special case where EXPR is an identifier goes a different path, through

WDL/__init__.py

@@ -77,6 +77,11 @@ def get_left(self, items, meta) -> E.Base:
     def get_right(self, items, meta) -> E.Base:
         return E.Apply(sp(self.filename, meta), "_get_right", items)
 
+    def map_kv(self, items, meta) -> E.Base:
+        assert len(items) == 2
+        return (items[0], items[1])
+    def map(self, items, meta) -> E.Base:
+        return E.Map(sp(self.filename, meta), items)
     def ifthenelse(self, items, meta) -> E.Base:
         assert len(items) == 3
         return E.IfThenElse(sp(self.filename, meta), *items)

WDL/_parser.py

@@ -46,6 +46,8 @@
           | expr_core "." _LEFT -> get_left
           | expr_core "." _RIGHT -> get_right
 
+          | "{" [map_kv ("," map_kv)*] "}" -> map
+
           | "if" expr "then" expr "else" expr -> ifthenelse
 
           | ident
@@ -80,6 +82,9 @@
 _RIGHT.2: "right"
 ident: [CNAME ("." CNAME)*]
 
+?map_key: literal | string
+map_kv: map_key ":" expr
+
 // WDL types and declarations
 type: _INT QUANT? -> int_type
     | _FLOAT QUANT? -> float_type

tests/test_eval.py

@@ -228,6 +228,15 @@ def test_pair(self):
             ("p.right", "2.71828", env)
         )
 
+    def test_map(self):
+        self._test_tuples(
+            ("{'foo': 1, 'bar': 2}['bar']", "2"),
+            ("{0: 1, 2: 3}['foo']", "", WDL.Error.StaticTypeMismatch),
+            ("{'foo': 1, 'bar': 2}[3]", "", WDL.Error.OutOfBounds), # int coerces to string...
+            ("{3: 1, false: 2}", "", WDL.Error.StaticTypeMismatch),
+            ("{'foo': true, 'bar': 0}", "", WDL.Error.StaticTypeMismatch)
+        )
+
     def test_errors(self):
         self._test_tuples(
             ("1 + bogus(2)", "(Ln 1, Col 5) No such function: bogus", WDL.Error.NoSuchFunction)