Merge pull request #16 from justbuchanan/matrix-init-and-access

implement Matrix.__getitem__() and Matrix(list) constructor

cadquery/occ_impl/geom.py

@@ -179,14 +179,43 @@ class Matrix:
     """A 3d , 4x4 transformation matrix.
 
     Used to move geometry in space.
+
+    The provided "matrix" parameter may be None, a gp_Trsf, or a nested list of
+    values.
+
+    If given a nested list, it is expected to be of the form:
+
+        [[m11, m12, m13, m14],
+         [m21, m22, m23, m24],
+         [m31, m32, m33, m34]]
+
+    A fourth row may be given, but it is expected to be: [0.0, 0.0, 0.0, 1.0]
+    since this is a transform matrix.
     """
 
     def __init__(self, matrix=None):
 
         if matrix is None:
             self.wrapped = gp_Trsf()
-        else:
+        elif isinstance(matrix, gp_Trsf):
             self.wrapped = matrix
+        elif isinstance(matrix, list):
+            self.wrapped = gp_Trsf()
+            if len(matrix) == 3:
+                flattened = [e for row in matrix for e in row]
+                self.wrapped.SetValues(*flattened)
+            elif len(matrix) == 4:
+                # Only use first 3 rows - the last must be [0, 0, 0, 1].
+                lastRow = matrix[3]
+                if lastRow != [0., 0., 0., 1.]:
+                    raise ValueError("Expected the last row to be [0,0,0,1], but got: {}".format(lastRow))
+                flattened = [e for row in matrix[0:3] for e in row]
+                self.wrapped.SetValues(*flattened)
+            else:
+                raise TypeError("Matrix constructor requires list of length 12 or 16")
+        else:
+            raise TypeError(
+                    "Invalid param to matrix constructor: {}".format(matrix))
 
     def rotateX(self, angle):
 
@@ -232,6 +261,25 @@ class Matrix:
         
         return [data[j][i] for i in range(4) for j in range(4)]
 
+    def __getitem__(self, rc):
+        """Provide Matrix[r, c] syntax for accessing individual values. The row
+        and column parameters start at zero, which is consistent with most
+        python libraries, but is counter to gp_Trsf(), which is 1-indexed.
+        """
+        if len(rc) != 2:
+            raise IndexError("Matrix subscript must provide (row, column)")
+        r, c = rc[0], rc[1]
+        if r >= 0 and r < 4 and c >= 0 and c < 4:
+            if r < 3:
+                return self.wrapped.Value(r+1,c+1)
+            else:
+                # gp_Trsf doesn't provide access to the 4th row because it has
+                # an implied value as below:
+                return [0., 0., 0., 1.][c]
+        else:
+            raise IndexError("Out of bounds access into 4x4 matrix: {}".format(rc))
+
+
 class Plane(object):
     """A 2D coordinate system in space
 

tests/TestCadObjects.py

@@ -134,6 +134,49 @@ class TestCadObjects(BaseTest):
         self.assertEqual(a, b)
         self.assertEqual(a, c)
 
+    def testMatrixCreationAndAccess(self):
+        def matrix_vals(m):
+            return [[m[r,c] for c in range(4)] for r in range(4)]
+        # default constructor creates a 4x4 identity matrix
+        m = Matrix()
+        identity = [[1., 0., 0., 0.],
+                    [0., 1., 0., 0.],
+                    [0., 0., 1., 0.],
+                    [0., 0., 0., 1.]]
+        self.assertEqual(identity, matrix_vals(m))
+
+        vals4x4 = [[1., 0., 0., 1.],
+                   [0., 1., 0., 2.],
+                   [0., 0., 1., 3.],
+                   [0., 0., 0., 1.]]
+
+        # test constructor with 16-value input
+        m = Matrix(vals4x4)
+        self.assertEqual(vals4x4, matrix_vals(m))
+
+        # test constructor with 12-value input (the last 4 are an implied
+        # [0,0,0,1])
+        m = Matrix(vals4x4[0:12])
+        self.assertEqual(vals4x4, matrix_vals(m))
+
+        # Test 16-value input with invalid values for the last 4
+        invalid = [[1., 0., 0., 1.],
+                   [0., 1., 0., 2.],
+                   [0., 0., 1., 3.],
+                   [1., 2., 3., 4.]]
+        with self.assertRaises(ValueError):
+            Matrix(invalid)
+
+        # Test input with invalid size
+        with self.assertRaises(TypeError):
+            Matrix([1,2,3])
+
+        # test out-of-bounds access
+        m = Matrix()
+        with self.assertRaises(IndexError):
+            m[5, 5]
+
+
     def testTranslate(self):
         e = Edge.makeCircle(2, (1, 2, 3))
         e2 = e.translate(Vector(0, 0, 1))