Types: handle MemType as argument to Array constructor

Alse add a to_type helper that factorizes the boilerplate code to
get a Type from either a MemType or a Type.

2 files changed, 29 insertions, 7 deletions

diff --git a/miasm2/core/types.py b/miasm2/core/types.py
index 2108d446..8ba45100 100644
--- a/miasm2/core/types.py
+++ b/miasm2/core/types.py
@@ -124,6 +124,29 @@ def set_allocator(alloc_func):
 
 # Helpers
 
+def to_type(obj):
+    """If possible, return the Type associated with @obj, otherwise raises
+    a ValueError.
+
+    Works with a Type instance (returns obj) or a MemType subclass or instance
+    (returns obj.get_type()).
+    """
+    # obj is a python type
+    if isinstance(obj, type):
+        if issubclass(obj, MemType):
+            if obj.get_type() is None:
+                raise ValueError("%r has no static type; use a subclasses "
+                                 "with a non null _type or use a "
+                                 "Type instance" % obj)
+            return obj.get_type()
+    # obj is not not a type
+    else:
+        if isinstance(obj, Type):
+            return obj
+        elif isinstance(obj, MemType):
+            return obj.get_type()
+    raise ValueError("%r is not a Type or a MemType" % obj)
+
 def indent(s, size=4):
     """Indent a string with @size spaces"""
     return ' '*size + ('\n' + ' '*size).join(s.split('\n'))
@@ -538,12 +561,7 @@ class Struct(Type):
         real_fields = []
         uniq_count = 0
         for fname, field in fields:
-            if isinstance(field, type) and issubclass(field, MemType):
-                if field._type is None:
-                    raise ValueError("%r has no static type; use a subclasses "
-                                     "with a non null _type or use a "
-                                     "Type instance")
-                field = field.get_type()
+            field = to_type(field)
 
             # For reflexion
             field._set_self_type(self)
@@ -726,7 +744,8 @@ class Array(Type):
     """
 
     def __init__(self, field_type, array_len=None):
-        self.field_type = field_type
+        # Handle both Type instance and MemType subclasses
+        self.field_type = to_type(field_type)
         self.array_len = array_len
 
     def _set_self_type(self, self_type):
diff --git a/test/core/test_types.py b/test/core/test_types.py
index bb1d5da1..09de1f71 100644
--- a/test/core/test_types.py
+++ b/test/core/test_types.py
@@ -154,6 +154,9 @@ assert memstr3.val == memstr.val # But the python value is the same
 
 
 # Array tests
+# Construction methods
+assert Array(MyStruct) == Array(MyStruct.get_type())
+assert Array(MyStruct, 10) == Array(MyStruct.get_type(), 10)
 # Allocate buffer manually, since memarray is unsized
 alloc_addr = my_heap.vm_alloc(jitter.vm, 0x100)
 memarray = Array(Num("I")).lval(jitter.vm, alloc_addr)