groupConnectivityNaive.py: doctest

groupConnectivityNaive.py

@@ -2,8 +2,25 @@ from sage.graphs.graph import Graph, DiGraph
 from sage.rings.finite_rings.integer_mod_ring import Integers
 
 def Zkn(k, n):
+  """Create group Z_k^n
+
+  EXAMPLES:
+
+  >>> list(Zkn(4, 1))
+  [0, 1, 2, 3]
+  >>> list(Zkn(3,2))
+  [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)]
+  >>> list(Zkn(2,3))  # doctest: +NORMALIZE_WHITESPACE
+  [(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1),
+   (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]
+  """
+
   Zk = Integers(k)
-  return Zk.cartesian_product( *([Zk] * (n - 1)) )
+  if n == 1:
+    return Zk
+  else:
+    return Zk.cartesian_product( *([Zk] * (n - 1)) )
+
 
 def prettifyGraph(G_):
   """Convert graph into directed one and number its edges."""
@@ -132,3 +149,9 @@ def myHistogram(data):
     s[i] += 1
   return s
 
+if __name__ == "__main__":
+    import doctest
+    ret = doctest.testmod()
+    print "%i tests of %i failed." % ret
+
+