Created som unit tests for statement class.

src/statement.py

@@ -5,7 +5,7 @@ class Statement:
     def __init__(self, stype, name, *args, **kwargs):
         self.stype = stype
         self.name = name
-        self.args = args
+        self.args = list(args)
         self.options = kwargs
 
     def __getitem__(self, n):
@@ -52,6 +52,12 @@ class Statement:
                and re.match('^j|jal|beq|bne|blez|bgtz|bltz|bgez|bct|bcf$', \
                             self.name)
 
+    def is_branch(self):
+        """Check if the statement is a branch."""
+        return self.is_command() \
+               and re.match('^beq|bne|blez|bgtz|bltz|bgez|bct|bcf$', \
+                            self.name)
+
     def is_shift(self):
         """Check if the statement is a shift operation."""
         return self.is_command() and re.match('^s(ll|la|rl|ra)$', self.name)
@@ -81,6 +87,9 @@ class Block:
     def __iter__(self):
         return iter(self.statements)
 
+    def __getitem__(self, n):
+        return self.statements[n]
+
     def __len__(self):
         return len(self.statements)
 

tests/test_statement.py

@@ -6,6 +6,7 @@ from src.statement import Statement as S, Block as B
 class TestStatement(unittest.TestCase):
 
     def setUp(self):
+        self.statement = S('command', 'foo', '$1')
         self.block = B([S('command', 'foo'), \
                         S('comment', 'bar'),
                         S('command', 'baz')])
@@ -13,6 +14,13 @@ class TestStatement(unittest.TestCase):
     def tearDown(self):
         del self.block
 
+    def test_getitem(self):
+        self.assertEqual(self.statement[0], '$1')
+
+    def test_setitem(self):
+        self.statement[0] = '$2'
+        self.assertEqual(self.statement[0], '$2')
+
     def test_eq(self):
         self.assertTrue(S('command', 'foo') == S('command', 'foo'))
         self.assertFalse(S('command', 'foo') == S('command', 'bar'))
@@ -69,3 +77,18 @@ class TestStatement(unittest.TestCase):
         self.block.apply_filter(lambda s: s.is_command())
         self.assertEqual(self.block.statements, [S('command', 'foo'), \
                                                  S('command', 'baz')])
+
+    def test_is_shift(self):
+        self.assertTrue(S('command', 'sll').is_shift())
+        self.assertFalse(S('command', 'foo').is_shift())
+        self.assertFalse(S('label', 'sll').is_shift())
+
+    def test_is_load(self):
+        self.assertTrue(S('command', 'lw').is_load())
+        self.assertFalse(S('command', 'foo').is_load())
+        self.assertFalse(S('label', 'lw').is_load())
+
+    def test_is_arith(self):
+        self.assertTrue(S('command', 'add', '$1', '$2', '$3').is_arith())
+        self.assertFalse(S('command', 'foo').is_arith())
+        self.assertFalse(S('label', 'add').is_arith())