diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..a60d365
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+
+*.un~
diff --git a/mpython.py b/mpython.py
index c8bfe70..0abba31 100755
--- a/mpython.py
+++ b/mpython.py
@@ -1,56 +1,67 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
-from ast import *
+import ast
 import sys
 
+
 def name(x, ctx):
-  return Name(id = x, ctx = ctx, lineno = 0, col_offset = 0)
+    return ast.Name(id=x, ctx=ctx, lineno=0, col_offset=0)
+
 
 def call(f, *args):
-  f = name(f, Load())
-  return Call(func = f, args = list(args), lineno = 0,
-              col_offset = 0, keywords = [], vararg = None)
+    f = name(f, ast.Load())
+    return ast.Call(func=f, args=list(args), lineno=0,
+                    col_offset=0, keywords=[], vararg=None)
+
 
 def func(args, body):
-  return Lambda(args = arguments(args = args, defaults = []),
-                body = body, vararg = None, lineno = 0, col_offset = 0)
+    return ast.Lambda(args=ast.arguments(args=args, defaults=[]),
+                      body=body, vararg=None, lineno=0, col_offset=0)
+
+
+unit = ast.Tuple(elts=[], lineno=0, col_offset=0, ctx=ast.Load())
 
-unit = Tuple(elts = [], lineno=0, col_offset = 0, ctx = Load())
 
 def transform(elt, generators):
-  elt = call("__singleton__", elt)
-
-  for generator in generators[-1::-1]:
-    for i in generator.ifs[-1::-1]:
-      elt = call("__concatMap__",
-                 func([name('_', Param())], elt),
-                 IfExp(i,
-                    call('__singleton__', unit),
-                    call('__fail__'), lineno=0, col_offset=0))
-
-    elt = call("__concatMap__",
-               func([name(generator.target.id, Param())], elt),
-               generator.iter)
-  return elt
-
-class RewriteComp(NodeTransformer):
-  def visit_ListComp(self, node):
-    return transform(RewriteComp().visit(node.elt),
-                     [RewriteComp().visit(generator) for generator in node.generators])
+    elt = call("__singleton__", elt)
+
+    for generator in reversed(generators):
+        for i in reversed(generator.ifs):
+            elt = call("__concatMap__",
+                       func([name('_', ast.Param())], elt),
+                       ast.IfExp(i,
+                                 call('__singleton__', unit),
+                                 call('__fail__'), lineno=0, col_offset=0))
+
+        elt = call("__concatMap__",
+                   func([name(generator.target.id, ast.Param())], elt),
+                   generator.iter)
+    return elt
+
+
+class RewriteComp(ast.NodeTransformer):
+    def visit_ListComp(self, node):
+        return transform(RewriteComp().visit(node.elt),
+                         [RewriteComp().visit(generator)
+                         for generator in node.generators])
+
 
 def __concatMap__(f, x):
-  return [z for y in x for z in f(y)]
+    return [z for y in x for z in f(y)]
+
 
 def __singleton__(x):
-  return [x]
+    return [x]
+
 
 def __fail__():
-  return []
+    return []
+
 
 source = open(sys.argv[1]).read()
-e = compile(source, "<string>", "exec", PyCF_ONLY_AST)
+e = compile(source, "<string>", "exec", ast.PyCF_ONLY_AST)
 e = RewriteComp().visit(e)
 f = compile(e, sys.argv[1], "exec")
 print f
 exec f
-print "Done"
\ No newline at end of file
+print "Done"
diff --git a/test1.py b/test1.py
index 01ff8e1..5f11a1f 100644
--- a/test1.py
+++ b/test1.py
@@ -1,21 +1,27 @@
 # Some simple regression tests
 
-a = [(x, y, z) for x in [0, 1, 2]
-               for y in [0, 1, 2]
-               for z in [0, 1, 2]]
+a = [
+    (x, y, z)
+    for x in [0, 1, 2]
+    for y in [0, 1, 2]
+    for z in [0, 1, 2]]
 
 print len(a)
 
-b = [(x, y) for x in [0,1]
-            for y in [0,1]
-            if x+y < 2]
+b = [
+    (x, y)
+    for x in [0, 1]
+    for y in [0, 1]
+    if x+y < 2]
 
 print b
 
-c = [(x, y, z) for x in range(1, 10)
-               for y in range(1, 10)
-               if x < y
-               for z in range(1, 10)
-               if x*x+y*y == z*z]
+c = [
+    (x, y, z)
+    for x in range(1, 10)
+    for y in range(1, 10)
+    if x < y
+    for z in range(1, 10)
+    if x*x+y*y == z*z]
 
-print c
\ No newline at end of file
+print c
diff --git a/test2.py b/test2.py
index 5ed6fca..2660f12 100644
--- a/test2.py
+++ b/test2.py
@@ -1,30 +1,39 @@
 import math
 
+
 def __concatMap__(k, m):
-  return lambda c:m(lambda a:k(a)(c))
+    return lambda c: m(lambda a: k(a)(c))
+
 
 def __singleton__(x):
-  return lambda f:f(x)
+    return lambda f: f(x)
+
 
 def callCC(f):
-  return lambda c:f(lambda a:lambda _:c(a))(c)
+    return lambda c: f(lambda a: lambda _: c(a))(c)
+
 
 def __fail__():
-  raise "Failure is not an option for continuations"
+    raise "Failure is not an option for continuations"
+
 
 def ret(x):
-  return __singleton__(x)
+    return __singleton__(x)
+
 
 def id(x):
-  return x
+    return x
+
 
 def solve(a, b, c):
-  return callCC(lambda throw: [((-b-d)/(2*a), (-b+d)/(2*a))
-                               for a0 in (ret(a) if a!=0 else throw("Not quadratic"))
-                               for d2 in ret(b*b-4*a*c)
-                               for d in (ret(math.sqrt(d2)) if d2>=0 else throw("No roots"))
-                              ])
+    return callCC(
+        lambda throw: [
+            ((-b-d)/(2*a), (-b+d)/(2*a))
+            for a0 in (ret(a) if a != 0 else throw("Not quadratic"))
+            for d2 in ret(b*b-4*a*c)
+            for d in (ret(math.sqrt(d2)) if d2 >= 0 else throw("No roots"))
+        ])
 
 print solve(1, 0, -9)(id)
 print solve(1, 1, 9)(id)
-print solve(0, 1, 9)(id)
\ No newline at end of file
+print solve(0, 1, 9)(id)
diff --git a/test3.py b/test3.py
index e8f7fdd..909dcc2 100644
--- a/test3.py
+++ b/test3.py
@@ -7,61 +7,72 @@
 # to mean something like "the probability of x supposing a is drawn from
 # distribution b and then c is drawn from d, conditioned on a
 #
+
+
 class PDF(dict):
     def __hash__(self):
         return hash(tuple(sorted(self.iteritems())))
 
+
 def scale(alpha, ps):
-  result = PDF()
-  for k in ps:
-    result[k] = alpha*ps[k]
-  return result
+    result = PDF()
+    for k in ps:
+        result[k] = alpha*ps[k]
+    return result
+
 
 def combine(pdf1, pdf2):
-  result  = pdf1.copy()
-  for k in pdf2.keys():
-    if result.has_key(k):
-      result[k] += pdf2[k]
-    else:
-      result[k] = pdf2[k]
-  return result
+    result = pdf1.copy()
+    for k in pdf2.keys():
+        if k in result:
+            result[k] += pdf2[k]
+        else:
+            result[k] = pdf2[k]
+    return result
+
 
 def __concatMap__(f, pps):
-  result = PDF()
-  for k0 in pps.keys():
-    p = pps[k0]
-    k = f(k0)
-    result = combine(result, scale(pps[k0], f(k0)))
-  return result
+    result = PDF()
+    for k0 in pps.keys():
+        p = pps[k0]
+        k = f(k0)
+        result = combine(result, scale(p, k))
+    return result
+
 
 def __singleton__(x):
-  return PDF([(x, 1)])
+    return PDF([(x, 1)])
+
 
 def __fail__(x):
-  return PDF()
+    return PDF()
+
 
 def certain(x):
-  return __singleton__(x)
+    return __singleton__(x)
+
 
 def expectation(pdf):
-  total = 0.0;
-  for k in pdf.keys():
-    total += pdf[k]*k
-  return total
+    total = 0.0
+    for k in pdf.keys():
+        total += pdf[k]*k
+    return total
+
 
 def probability(condition, pdf):
-  return expectation([1 if condition(k) else 0 for k in pdf])
+    return expectation([1 if condition(k) else 0 for k in pdf])
 
-################################################################################
+###############################################################################
 # Test code starts here
-################################################################################
+###############################################################################
+
 
-def bernoulli(p, a = True, b = False):
-  """
-  A sample from bernoulli(p, a, b) has probability p of taking the value
-  a and probabilty 1-p of taking the value b.
-  """
-  return PDF([(a, p), (b, 1-p)])
+def bernoulli(p, a=True, b=False):
+    """
+    A sample from bernoulli(p, a, b) has probability p of taking the value
+    a and probabilty 1-p of taking the value b.
+    """
+    return PDF([(a, p), (b, 1-p)])
 
 # Straightforward coin toss.
 test1 = [x for x in bernoulli(0.5)]
@@ -69,53 +80,64 @@ def bernoulli(p, a = True, b = False):
 print "test1 =", test1
 
 # A pair of coin tosses.
-test2 = [(x, y) for x in bernoulli(0.5)
-                for y in bernoulli(0.5)]
+test2 = [
+    (x, y)
+    for x in bernoulli(0.5)
+    for y in bernoulli(0.5)]
 
 print "test2 =", test2
 
 # The sum of the result of three coin tosses.
-test3 = [a+b+c for a in bernoulli(0.5, 1, 0)
-               for b in bernoulli(0.5, 1, 0)
-               for c in bernoulli(0.5, 1, 0)]
+test3 = [
+    a+b+c
+    for a in bernoulli(0.5, 1, 0)
+    for b in bernoulli(0.5, 1, 0)
+    for c in bernoulli(0.5, 1, 0)]
 
 print "test3 =", test3
 
-################################################################################
+###############################################################################
 # Chinese restaurant code
 # http://en.wikipedia.org/wiki/Chinese_restaurant_process
-################################################################################
+###############################################################################
+
 
 def addGuest(i, tables):
-  if i==len(tables):
-    return tables+(1,)
-  else:
-    tableList = list(tables)
-    tableList[i] += 1
+    if i == len(tables):
+        return tables+(1,)
+    else:
+        tableList = list(tables)
+        tableList[i] += 1
     return tuple(tableList)
 
+
 def selectTable(tables):
-  newTables = list(tables)+[1]
-  n = sum(newTables)
-  return PDF(zip(range(len(newTables)), map(lambda x: float(x)/n, newTables)))
+    newTables = list(tables)+[1]
+    n = sum(newTables)
+    return PDF(
+        zip(range(len(newTables)), map(lambda x: float(x)/n, newTables)))
+
 
 def chineseRestaurantProcess(n):
-  """
-  Simulates the Chinese restaurant process to n steps
-  """
+    """
+    Simulates the Chinese restaurant process to n steps
+    """
+
+    if n == 0:
+        return certain(())
+    else:
+        return [
+            addGuest(guestTable, tables)
+            for tables in chineseRestaurantProcess(n-1)
+            for guestTable in selectTable(tables)]
 
-  if n==0:
-    return certain(())
-  else:
-    return [addGuest(guestTable, tables) for tables in chineseRestaurantProcess(n-1)
-                                         for guestTable in selectTable(tables)]
 
 # Expected table size (occupancy) given by theory.
 def theoreticalTableSize(n):
-  total = 0.0
-  for k in range(0, n):
-    total += 1.0/(k+1)
-  return total
+    total = 0.0
+    for k in range(0, n):
+        total += 1.0/(k+1)
+    return total
 
 # Compare theoretical and simulated expected table size.
 # Should be equal to machine precision.
@@ -128,18 +150,23 @@ def theoreticalTableSize(n):
 # is probably much more accurate than a Monte-Carlo simulation using the same
 # amount of CPU time.
 for i in range(1, 16):
-  test5a = probability(lambda table: min(table)==1, chineseRestaurantProcess(i))
-  print i, test5a
+    test5a = probability(
+        lambda table: min(table) == 1, chineseRestaurantProcess(i))
+    print i, test5a
 test5b = 1-math.exp(-1)
 print "test5 =", math.fabs(test5a-test5b) < 1e-9
 
 # Maybe you recognise the 1-1/e above from this problem:
-# n people visit a restaurant (another one) and hang up their hats. When they leave
-# they pick a hat at random. What's the probability of anyone getting their original
+# n people visit a restaurant (another one) and hang up their hats.
+# When they leave they pick a hat at random.
+# What's the probability of anyone getting their original
 # hat back? The connection between that problem and the CRP is a beautiful bit
-# of mathematics about the permutation group. I'll leave it as a puzzle to figure
+# of mathematics about the permutation group.
+# I'll leave it as a puzzle to figure
 # out why problems have the same solution.
 
-# PS I only realised this fact because I'd written code to compute this stuff to machine
-# precision and recognised 1-1/e. It's fortuitous that I chose this problem completely at
-# random and was led to see the connection. (I admit I got some help from the web...)
+# PS I only realised this fact because I'd written code to compute
+# this stuff to machine precision and recognised 1-1/e.
+# It's fortuitous that I chose this problem completely at
+# random and was led to see the connection.
+# (I admit I got some help from the web...)
diff --git a/test4.py b/test4.py
new file mode 100644
index 0000000..9fb5c22
--- /dev/null
+++ b/test4.py
@@ -0,0 +1,21 @@
+# Maybe Monad
+
+def __concatMap__(f, m):
+    return None if m is None else f(m)
+
+def __singleton__(x):
+    return x
+
+def __fail__():
+    return None
+
+def ret(x):
+    return __singleton__(x)
+
+print [
+    x
+    for x in 242
+    for y in x + 30
+    for z in y + 40
+    for q in None]
+