Adding AugAssign List String

pyObjectManager/List.py

@@ -140,7 +140,7 @@ def __setitem__(self,key,value):
         """
         # Attempt to return variable
         assert type(key) is int
-        assert type(value) in [Int, Real, BitVec, List]
+        assert type(value) in [Int, Real, BitVec, List, String]
 
         # Get that index's current count
         count = self.variables[key].count + 1
@@ -160,10 +160,10 @@ def __setitem__(self,key,value):
             self.variables[key] = BitVec('{2}{0}[{1}]'.format(self.varName,key,self.count),ctx=self.ctx,count=count,size=value.size)
             self.state.addConstraint(self.variables[key].getZ3Object() == value.getZ3Object())
 
-        elif type(value) is List:
-            logger.debug("__setitem__: setting List")
+        elif type(value) in [List, String]:
+            logger.debug("__setitem__: setting {0}".format(type(value)))
             self.variables[key] = value
-            value.count = count
+            #value.count = count
 
         else:
             err = "__setitem__: Don't know how to set object '{0}'".format(value)

pyState/AugAssign.py

@@ -6,7 +6,7 @@
 from pyObjectManager.Int import Int
 from pyObjectManager.Real import Real
 from pyObjectManager.BitVec import BitVec
-
+from pyObjectManager.String import String
 
 logger = logging.getLogger("pyState:AugAssign")
 
@@ -99,6 +99,32 @@ def _handleNum(state,element,value,op):
     # Return the state
     return [state]
 
+def _handleString(state,element,value,op):
+    """
+    Handle the case where we're AugAssigning Strings
+    """
+
+    # Find the parent object
+    oldTarget = state.resolveObject(element.target)
+    parent = state.objectManager.getParent(oldTarget)
+    index = parent.index(oldTarget)
+
+    # Create a new string
+    newString = state.getVar("AugAssignTempString",ctx=1,varType=String)
+    newString.increment()
+
+    # Set the new string
+    newString.variables = oldTarget.variables + value.variables
+
+    # Assign the new string
+    parent[index] = newString.copy()
+
+    # Pop the instruction off
+    state.path.pop(0)
+
+    # Return the state
+    return [state]
+
 
 def handle(state,element):
     """
@@ -125,5 +151,13 @@ def handle(state,element):
 
         if type(value) in [Int, BitVec, Real]:
             ret += _handleNum(state.copy(),element,value,op)
-
+
+        elif type(value) is String:
+            ret += _handleString(state.copy(),element,value,op)
+
+        else:
+            err = "handle: Don't know how to handle type {0}".format(type(value))
+            logger.error(err)
+            raise Exception(err)
+
     return ret

tests/test_pyObjectManager_List.py

@@ -99,7 +99,7 @@ def test_pyObjectManager_List_setitem():
 
     # Assign List
     l[1] = List(varName='x',ctx=0)
-    assert l[1].count == 4
+    #assert l[1].count == 4
     assert type(l[1]) == List
 
 

tests/test_pyState_AugAssign.py

@@ -97,6 +97,24 @@ def test(x,y):
 x += s.index("b")
 """
 
+test14 = """
+l = ["a","b","c"]
+l[1] += "d"
+l[1] += "E"
+"""
+
+def test_pySym_AugAssign_String_In_List():
+    b = ast.parse(test14).body
+    p = Path(b,source=test14)
+    pg = PathGroup(p)
+    pg.explore()
+
+    # There should be 10 possible states here
+    assert len(pg.completed) == 1
+
+    assert pg.completed[0].state.any_list('l') == ['a', 'bdE', 'c']
+
+
 def test_pySym_AugAssign_MultipleStates():
     b = ast.parse(test13).body
     p = Path(b,source=test13)