Recognize streaming shapes correctly

youtype · Jun 6, 2020 · 1ec25dc · 1ec25dc
1 parent c3e6aca
commit 1ec25dc
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/mypy_boto3_builder/parsers/shape_parser.py b/mypy_boto3_builder/parsers/shape_parser.py
@@ -60,7 +60,6 @@ class ShapeParser:
         "double": Type.float,
         "float": Type.float,
         "timestamp": ExternalImport(ImportString("datetime"), "datetime"),
-        "blob": TypeSubscript(Type.Union, [Type.bytes, Type.IOBytes]),
     }
 
     # Alias map fixes added by botocore for documentation build.
@@ -319,6 +318,11 @@ def _parse_shape_list(self, shape: ListShape) -> FakeAnnotation:
         return type_subscript
 
     def _parse_shape(self, shape: Shape) -> FakeAnnotation:
+        if shape.type_name == "blob":
+            if shape.serialization.get("streaming"):
+                return Type.IOBytes
+            return Type.bytes
+
         if shape.type_name in self.SHAPE_TYPE_MAP:
             return self.SHAPE_TYPE_MAP[shape.type_name]