Update TypeRegistry

Author: jterapin

gems/smithy-schema/lib/smithy-schema/document.rb

@@ -35,7 +35,8 @@ def [](key)
         @data[key]
       end
 
-      # @param [Smithy::Schema::Shapes::StructureShape] schema
+      # @param [Shapes::Shape] schema
+      # @return [Shapes::Structure] typed shape
       def as_typed(schema)
         error_message = 'Invalid schema or document data'
         raise ArgumentError, error_message unless valid_schema?(schema) && @data.is_a?(Hash)

gems/smithy-schema/lib/smithy-schema/type_registry.rb

@@ -2,14 +2,17 @@
 
 module Smithy
   module Schema
-    # Registry that contains a map of Smithy shape ID to its shape representation
-    # TODO: Implement a method that takes a document and deserializes
-    #  Do we include operation shapes in this registry?
+    # A registry that contains a map of Smithy shape ID to its schema.
+    # Also includes a way to find schema based on its shape type representation
     class TypeRegistry
       def initialize
         @registry = {}
+        @schema_by_types = {}
       end
 
+      # @api private
+      attr_reader :schema_by_types
+
       # @return [Hash<String, Shapes::Shape>]
       attr_accessor :registry
 
@@ -18,28 +21,62 @@ def register(*shapes)
         raise ArgumentError, 'Expected an array of Shapes' unless shapes.all?(Shapes::Shape)
 
         shapes.each do |s|
-          next if s.id.nil?
-
           @registry[s.id] = s
+
+          case s.class
+          when Shapes::StructureShape
+            @schema_by_types[s.type] = s if s.type
+          when Shapes::UnionShape
+            s.member_types.values { |v| @schema_by_types[v] = s }
+          end
         end
       end
 
-      # @param [String] shape_id
+      # Returns true if this type registry contains specific shape id.
+      # @param [String] id
+      # @return [Boolean]
+      def schema_by_id?(id)
+        @registry.key?(id)
+      end
+
+      # Returns the shape schema registered for the given shape id.
+      # @param [id] id
+      # @return [Shapes::Shape, nil]
+      def schema_by_id(id)
+        @registry[id]
+      end
+
+      # Returns true if this type registry contains a schema associated
+      #  with the given typed shape.
+      # @param [Class] type
       # @return [Boolean]
-      def shape?(shape_id)
-        @registry.key?(shape_id)
+      def scheme_by_type?(type)
+        @schema_by_types.key?(type)
+      end
+
+      # Returns the shape schema registered for the given typed shape.
+      # @param [Class] type
+      # @return [Shapes::Shape, nil]
+      def scheme_by_type(type)
+        @schema_by_types[type]
       end
 
-      # @param [String] shape_id
-      # @return [Shapes::Shape]
-      def shape(shape_id)
-        @registry[shape_id]
+      # Deserializes a document into a typed shape from registry.
+      # @param [Document] document
+      # @return [Shapes::Structure] typed shape
+      def convert_as_typed(document)
+        msg = 'Unable to convert given document since discriminator is not set'
+        raise ArgumentError, msg unless document.discriminator
+
+        if schema_by_id?(document.discriminator)
+          document.as_typed(schema_by_id(document.discriminator))
+        else
+          msg = "Unable to find schema with #{document.discriminator} in Registry"
+          raise ArgumentError, msg
+        end
       end
 
       class << self
-        # TODO: Need thoughts on...
-        #  * Smithy-Java only allows to compose 2 type registries at a time.
-        #    Do we follow suit or allow unlimited number of registry to compose?
         # @param [Array<TypeRegistry>]
         # @return [TypeRegistry]
         def compose(*type_registries)