MemStruct: allow Type instance in cast

3 files changed, 19 insertions, 16 deletions

diff --git a/example/jitter/memstruct.py b/example/jitter/memstruct.py
index 5472798d..3b6358cd 100644
--- a/example/jitter/memstruct.py
+++ b/example/jitter/memstruct.py
@@ -214,10 +214,10 @@ print
 raw_miams = '\x00'.join('Miams') + '\x00'*3
 raw_miams_array = [ord(c) for c in raw_miams]
 assert list(data.array)[:len(raw_miams_array)] == raw_miams_array
-assert data.array.cast(Str("utf16").pinned) == memstr
+assert data.array.cast(Str("utf16")) == memstr
 # Default is "ansi"
-assert data.array.cast(Str().pinned) != memstr
-assert data.array.cast(Str("utf16").pinned).val == memstr.val
+assert data.array.cast(Str()) != memstr
+assert data.array.cast(Str("utf16")).val == memstr.val
 
 print "See that the original array has been modified:"
 print repr(data)
diff --git a/miasm2/analysis/mem.py b/miasm2/analysis/mem.py
index 9787a25e..4d9ac712 100644
--- a/miasm2/analysis/mem.py
+++ b/miasm2/analysis/mem.py
@@ -946,11 +946,16 @@ class PinnedType(object):
             raise ValueError("byte must be a 1-lengthed str")
         self._vm.set_mem(self.get_addr(), byte * self.get_size())
 
-    def cast(self, other_type, *type_args, **type_kwargs):
-        """Cast this PinnedType to another PinnedType (same address, same vm, but
-        different type). Return the casted PinnedType.
+    def cast(self, other_type):
+        """Cast this PinnedType to another PinnedType (same address, same vm,
+        but different type). Return the casted PinnedType.
+
+        @other_type: either a Type instance (other_type.pinned is used) or a
+            PinnedType subclass
         """
-        return other_type(self._vm, self.get_addr(), *type_args, **type_kwargs)
+        if isinstance(other_type, Type):
+            other_type = other_type.pinned
+        return other_type(self._vm, self.get_addr())
 
     def cast_field(self, field, other_type, *type_args, **type_kwargs):
         """ABSTRACT: Same as cast, but the address of the returned PinnedType
@@ -958,6 +963,8 @@ class PinnedType(object):
 
         @field: field specification, for example its name for a struct, or an
             index in an array. See the subclass doc.
+        @other_type: either a Type instance (other_type.pinned is used) or a
+            PinnedType subclass
         """
         raise NotImplementedError("Abstract")
 
@@ -1074,12 +1081,13 @@ class PinnedStruct(PinnedType):
         """
         return self._type.set_field(self._vm, self.get_addr(), name, val)
 
-    def cast_field(self, field, other_type, *type_args, **type_kwargs):
+    def cast_field(self, field, other_type):
         """
         @field: a field name
         """
-        return other_type(self._vm, self.get_addr(field),
-                          *type_args, **type_kwargs)
+        if isinstance(other_type, Type):
+            other_type = other_type.pinned
+        return other_type(self._vm, self.get_addr(field))
 
 
     # Field generation methods, voluntarily public to be able to gen fields
@@ -1277,10 +1285,6 @@ class PinnedArray(PinnedType):
     def __setitem__(self, idx, item):
         self.get_type().set_item(self._vm, self._addr, idx, item)
 
-    # just a shorthand
-    def as_mem_str(self, encoding="ansi"):
-        return self.cast(Str(encoding).pinned)
-
     def raw(self):
         raise ValueError("%s is unsized, which prevents from getting its full "
                          "raw representation. Use PinnedSizedArray instead." %
diff --git a/test/analysis/mem.py b/test/analysis/mem.py
index 8d4a56d3..90022fe9 100644
--- a/test/analysis/mem.py
+++ b/test/analysis/mem.py
@@ -426,11 +426,10 @@ assert PinnedShort(jitter.vm, ms2.s2.get_addr(4)).val == 0xabcd
 
 # void* style cast
 PinnedPtrVoid = Ptr("I", Void()).pinned
-PinnedPtrMyStruct = Ptr("I", MyStruct).pinned
 p = PinnedPtrVoid(jitter.vm)
 p.val = mstruct.get_addr()
 assert p.deref.cast(MyStruct) == mstruct
-assert p.cast(PinnedPtrMyStruct).deref == mstruct
+assert p.cast(Ptr("I", MyStruct)).deref == mstruct
 
 # Field equality tests
 assert RawStruct("IH") == RawStruct("IH")