Generate more compact switch statements for polymorphic types.

[apitrace] / specs / stdapi.py

diff --git a/specs/stdapi.py b/specs/stdapi.py
index 1b9bf58ad5f5ab0f4b6cb57f5f6f275048e3f628..fa371cf2d8706d8286d88d9d9ee0065f18123eed 100644 (file)
--- a/specs/stdapi.py
+++ b/specs/stdapi.py
@@ -349,6 +349,22 @@ class Polymorphic(Type):
     def visit(self, visitor, *args, **kwargs):
         return visitor.visit_polymorphic(self, *args, **kwargs)
 
+    def iterswitch(self):
+        cases = [['default']]
+        types = [self.default_type]
+
+        for expr, type in self.switch_types:
+            case = 'case %s' % expr
+            try:
+                i = types.index(type)
+            except ValueError:
+                cases.append([case])
+                types.append(type)
+            else:
+                cases[i].append(case)
+
+        return zip(cases, types)
+
 
 class Visitor: