Added dictionary comprehension

Author: dmitrii-eremin

pythonlua/luainit.lua

@@ -48,6 +48,13 @@ end
 
 local function list(t)
     local result = {}
+
+    result.is_list = true
+
+    result._data = {}
+    for _, v in ipairs(t) do
+        table.insert(result._data, v)
+    end
 
     local methods = {}
 
@@ -65,6 +72,9 @@ local function list(t)
 
             return methods[index]
         end,
+        __newindex = function(self, index, value)
+            result._data[index] = value
+        end,
         __call = function(self, _, idx)
             if idx == nil and iterator_index ~= nil then
                 iterator_index = nil
@@ -77,19 +87,19 @@ local function list(t)
         end,
     })
 
-    result.is_list = true
-
-    result._data = {}
-    for _, v in ipairs(t) do
-        table.insert(result._data, v)
-    end
-
     return result
 end
 
 function dict(t)
     local result = {}
 
+    result.is_dict = true
+
+    result._data = {}
+    for k, v in pairs(t) do
+        result._data[k] = v
+    end
+
     local methods = {}
 
     methods.items = function()
@@ -98,8 +108,13 @@ function dict(t)
 
     local key_index = nil
 
-    setmetatable(t, {
-        __index = methods,
+    setmetatable(result, {
+        __index = function(self, index)
+            return methods[index]
+        end,
+        __newindex = function(self, index, value)
+            result._data[index] = value
+        end,
         __call = function(self, _, idx)
             if idx == nil and key_index ~= nil then
                 key_index = nil
@@ -110,13 +125,6 @@ function dict(t)
             return key_index
         end,
     })
-
-    result.is_dict = true
-
-    result._data = {}
-    for k, v in pairs(t) do
-        result._data[k] = v
-    end
 
     return result
 end

pythonlua/nodevisitor.py

@@ -192,6 +192,40 @@ def visit_Dict(self, node):
         elements = ", ".join(elements)
         self.emit("dict {{{}}}".format(elements))
 
+    def visit_DictComp(self, node):
+        """Visit dictionary comprehension"""
+        self.emit("(function()")
+        self.emit("local result = dict {}")
+
+        ends_count = 0
+
+        for comp in node.generators:
+            line = "for {target} in {iterator} do"
+            values = {
+                "target": self.visit_all(comp.target, inline=True),
+                "iterator": self.visit_all(comp.iter, inline=True),
+            }
+            line = line.format(**values)
+            self.emit(line)
+            ends_count += 1
+
+            for if_ in comp.ifs:
+                line = "if {} then".format(self.visit_all(if_, inline=True))
+                self.emit(line)
+                ends_count += 1
+
+        line = "result[{key}] = {value}"
+        values = {
+            "key": self.visit_all(node.key, inline=True),
+            "value": self.visit_all(node.value, inline=True),
+        }
+        self.emit(line.format(**values))
+
+        self.emit(" ".join(["end"] * ends_count))
+
+        self.emit("return result")
+        self.emit("end)()")
+
     def visit_Ellipsis(self, node):
         """Visit ellipsis"""
         self.emit("...")

tests/comprehensions.py

@@ -1,4 +1,9 @@
 a = [i * j for i in range(5) for j in range(3) if i * j % 2 == 0 and i > 0 and j > 0]
 
 for item in a:
-    print(item)
+    print(item)
+
+b = {i: i ** 2 for i in range(5)}
+
+for k, v in b.items():
+    print(k, v)

tests/comprehensions.py.expected

@@ -7,3 +7,8 @@
 8
 12
 10
+1	1.0
+2	4.0
+3	9.0
+4	16.0
+5	25.0