From 0d69ed46bd77a9940309dde25d24cd2c5ce3d51b Mon Sep 17 00:00:00 2001
From: Sergiu Ivanov <unlimitedscolobb@gmail.com>
Date: Sun, 22 Jul 2012 17:11:23 +0300
Subject: [PATCH] DiagramGrid: Support any iterable to describe groups.

Previously DiagramGrid would only agree if groups were supplied as
FiniteSet's of FiniteSet's.  This commit enables DiagramGrid to cope
with any iterable type.  The test added hereby shows how this can be
utilised.
---
 sympy/categories/diagram_drawing.py    | 26 +++++++++++++++++++++++++-
 sympy/categories/tests/test_drawing.py | 26 ++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/sympy/categories/diagram_drawing.py b/sympy/categories/diagram_drawing.py
index 4c686cf02a03..57e99b365049 100644
--- a/sympy/categories/diagram_drawing.py
+++ b/sympy/categories/diagram_drawing.py
@@ -87,6 +87,7 @@
                               NamedMorphism, Diagram)
 from sympy.utilities import default_sort_key
 from itertools import chain
+from sympy.core.compatibility import iterable
 
 class _GrowableGrid(object):
     """
@@ -787,15 +788,38 @@ def lay_out_group(group, local_hints):
             else:
                 obj_groups[group] = group
 
+        def group_to_finiteset(group):
+            """
+            Converts ``group`` to a :class:``FiniteSet`` if it is an
+            iterable.
+            """
+            if iterable(group):
+                return FiniteSet(group)
+            else:
+                return group
+
         obj_groups = {}
         groups_grids = {}
 
+        # We would like to support various containers to represent
+        # groups.  To achieve that, before laying each group out, it
+        # should be converted to a FiniteSet, because that is what the
+        # following code expects.
+
         if isinstance(groups, dict) or isinstance(groups, Dict):
+            finiteset_groups = {}
             for group, local_hints in groups.items():
+                finiteset_group = group_to_finiteset(group)
+                finiteset_groups[finiteset_group] = local_hints
                 lay_out_group(group, local_hints)
+            groups = finiteset_groups
         else:
+            finiteset_groups = []
             for group in groups:
-                lay_out_group(group, None)
+                finiteset_group = group_to_finiteset(group)
+                finiteset_groups.append(finiteset_group)
+                lay_out_group(finiteset_group, None)
+            groups = finiteset_groups
 
         new_morphisms = []
         for morphism in merged_morphisms:
diff --git a/sympy/categories/tests/test_drawing.py b/sympy/categories/tests/test_drawing.py
index 650dd7f03e9e..ec2f6dddbe4f 100644
--- a/sympy/categories/tests/test_drawing.py
+++ b/sympy/categories/tests/test_drawing.py
@@ -446,6 +446,32 @@ def test_DiagramGrid():
     assert grid[2, 5] == E
     assert grid.morphisms == morphisms
 
+    # Test the five lemma with object grouping, but mixing containers
+    # to represent groups.
+    grid = DiagramGrid(d, [(A, B, C, D, E), set([A_, B_, C_, D_, E_])])
+
+    assert grid.width == 6
+    assert grid.height == 3
+    assert grid[0, 0] == A_
+    assert grid[0, 1] == B_
+    assert grid[0, 2] is None
+    assert grid[0, 3] == A
+    assert grid[0, 4] == B
+    assert grid[0, 5] is None
+    assert grid[1, 0] is None
+    assert grid[1, 1] == C_
+    assert grid[1, 2] == D_
+    assert grid[1, 3] is None
+    assert grid[1, 4] == C
+    assert grid[1, 5] == D
+    assert grid[2, 0] is None
+    assert grid[2, 1] is None
+    assert grid[2, 2] == E_
+    assert grid[2, 3] is None
+    assert grid[2, 4] is None
+    assert grid[2, 5] == E
+    assert grid.morphisms == morphisms
+
     # Test the five lemma with object grouping and hints.
     grid = DiagramGrid(d, {
         FiniteSet(A, B, C, D, E): {"layout": "sequential",