Implement Typed Documents and TypeRegistry

README.md

@@ -27,7 +27,7 @@ bundle exec smithy-ruby smith client --gem-name weather --gem-version 1.0.0 --de
 ### IRB
 IRB on `weather` gem:
 ```
-irb -I projections/weather/lib -I gems/smithy-client/lib -I gems/smithy-schema/lib -r weather
+irb -I projections/weather/lib -I gems/smithy-client/lib -I gems/smithy-schema/lib -I gems/smithy-cbor/lib -r weather
 ```
 
 Create a Weather client:

gems/smithy-client/lib/smithy-client/base.rb

@@ -161,6 +161,16 @@ def service=(service)
           define_operation_methods
         end
 
+        # @return [Schema::TypeRegistry]
+        def type_registry
+          @type_registry ||= Schema::TypeRegistry.new
+        end
+
+        # @param [TypeRegistry] registry
+        def set_type_registry=(registry)
+          @type_registry = registry
+        end
+
         # @option options [ServiceShape] :service (ServiceShape.new)
         # @option options [Array<Plugin>] :plugins ([]) A list of plugins to
         #  add to the client class created.

gems/smithy-client/spec/smithy-client/base_spec.rb

@@ -251,6 +251,20 @@ module Client
         end
       end
 
+      describe '.type_registry' do
+        it 'defaults to a TypeRegistry' do
+          expect(client_class.type_registry).to be_kind_of(Schema::TypeRegistry)
+        end
+      end
+
+      describe '.set_type_registry' do
+        it 'can be set' do
+          registry = Schema::TypeRegistry.new
+          client_class.set_type_registry = registry
+          expect(client_class.type_registry).to be(registry)
+        end
+      end
+
       describe '.define' do
         it 'creates a new client class' do
           client_class = Base.define

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

@@ -2,6 +2,9 @@
 
 require_relative 'smithy-schema/shapes'
 require_relative 'smithy-schema/structure'
+require_relative 'smithy-schema/document'
+require_relative 'smithy-schema/time_helper'
+require_relative 'smithy-schema/type_registry'
 require_relative 'smithy-schema/union'
 
 module Smithy

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

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+require_relative 'document_utils'
+
+module Smithy
+  module Schema
+    # A Smithy document type, representing typed or untyped data from Smithy data model.
+    # ## Document types
+    # Document types are protocol-agnostic view of untyped data. They could be combined
+    # with a shape to serialize its contents.
+    #
+    # Smithy-Ruby currently only support JSON documents.
+    class Document
+      # @param  [Object] data  document data
+      # @param [Hash] options
+      # @option options [Smithy::Schema::Structure] :shape shape to reference when setting
+      #  document data. Only applicable when data param is a type of {Shapes::StructureShape}.
+      # @option options [Boolean] :use_timestamp_format Whether to use the `timestampFormat`
+      #  trait or ignore it when creating a {Document} with given shape. The `timestampFormat`
+      #  trait is ignored by default.
+      # @option options [Boolean] :use_json_name Whether to use the `jsonName` trait or ignore
+      #  it when creating a {Document} with given shape. The `jsonName` trait is ignored
+      #  by default.
+      def initialize(data, options = {})
+        @data = set_data(data, options)
+        @discriminator = extract_discriminator(data, options)
+      end
+
+      # @return [Object] data
+      attr_reader :data
+
+      # @return [String] discriminator
+      attr_reader :discriminator
+
+      # @param [String] key
+      # @return [Object]
+      def [](key)
+        return unless @data.is_a?(Hash) && @data.key?(key)
+
+        @data[key]
+      end
+
+      # @param [Shapes::Structure] shape
+      # @return [Object] typed shape
+      def as_typed(shape)
+        error_message = 'Invalid shape or document data'
+        raise ArgumentError, error_message unless valid_shape?(shape) && @data.is_a?(Hash)
+
+        type = shape.type.new
+        DocumentUtils.apply(@data, shape, type)
+      end
+
+      private
+
+      def discriminator?(data)
+        data.is_a?(Hash) && data.key?('__type')
+      end
+
+      def extract_discriminator(data, opts)
+        return if data.nil?
+
+        return unless discriminator?(data) || (shape = opts[:shape])
+
+        if discriminator?(data)
+          data['__type']
+        else
+          error_message = "Expected a structure shape, given #{shape.class} instead"
+          raise error_message unless valid_shape?(shape)
+
+          shape.id
+        end
+      end
+
+      def set_data(data, opts)
+        return if data.nil?
+
+        case data
+        when Smithy::Schema::Structure
+          shape = opts[:shape]
+          if shape.nil? || !valid_shape?(shape)
+            raise ArgumentError, "Unable to create a document with given shape: #{shape}"
+          end
+
+          opts = opts.except(:shape)
+          DocumentUtils.extract(data, shape, opts)
+        else
+          if discriminator?(data)
+            data.except('__type')
+          else
+            DocumentUtils.format(data)
+          end
+        end
+      end
+
+      def valid_shape?(shape)
+        shape.is_a?(Shapes::StructureShape) && !shape.type.nil?
+      end
+    end
+  end
+end

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

@@ -0,0 +1,217 @@
+# frozen_string_literal: true
+
+require 'base64'
+require 'time'
+
+module Smithy
+  module Schema
+    # @api private
+    # Document Utilities to help (de)construct data to/from Smithy document
+    module DocumentUtils
+      class << self
+        # Used to transform untyped data
+        def format(data)
+          return if data.nil?
+
+          case data
+          when Time
+            data.to_i # timestamp format is "epoch-seconds" by default
+          when Hash
+            data.each_with_object({}) do |(k, v), h|
+              h[k.to_s] = format(v)
+            end
+          when Array
+            data.map { |d| format(d) }
+          else
+            data
+          end
+        end
+
+        # Used to apply data to runtime shape
+        def apply(data, shape, type = nil)
+          case shape_reference(shape)
+          when Shapes::StructureShape then apply_structure(data, shape, type)
+          when Shapes::UnionShape then apply_union(data, shape, type)
+          when Shapes::ListShape then apply_list(data, shape)
+          when Shapes::MapShape then apply_map(data, shape)
+          when Shapes::TimestampShape then apply_timestamp(data, shape)
+          when Shapes::BlobShape then Base64.decode64(data)
+          else data
+          end
+        end
+
+        # rubocop:disable Metrics/CyclomaticComplexity
+        def extract(data, shape, opts = {})
+          return if data.nil?
+
+          case shape_reference(shape)
+          when Shapes::StructureShape then extract_structure(data, shape, opts)
+          when Shapes::UnionShape     then extract_union(data, shape, opts)
+          when Shapes::ListShape      then extract_list(data, shape)
+          when Shapes::MapShape       then extract_map(data, shape)
+          when Shapes::BlobShape      then extract_blob(data)
+          when Shapes::TimestampShape then extract_timestamp(data, shape, opts)
+          else data
+          end
+        end
+        # rubocop:enable Metrics/CyclomaticComplexity
+
+        private
+
+        def apply_list(data, shape)
+          shape = shape_reference(shape)
+          data.map do |v|
+            next if v.nil?
+
+            apply(v, shape.member)
+          end
+        end
+
+        def apply_map(data, shape)
+          shape = shape_reference(shape)
+          data.transform_values do |v|
+            if v.nil?
+              nil
+            else
+              apply(v, shape.value)
+            end
+          end
+        end
+
+        def apply_structure(data, shape, type)
+          shape = shape_reference(shape)
+
+          type = shape.type.new if type.nil?
+          data.each do |k, v|
+            name =
+              if (member = member_with_json_name(k, shape))
+                shape.name_by_member_name(member.name)
+              else
+                member_name(shape, k)
+              end
+            next if name.nil?
+
+            type[name] = apply(v, shape.member(name))
+          end
+          type
+        end
+
+        def apply_timestamp(data, shape)
+          data = data.is_a?(Numeric) ? Time.at(data) : Time.parse(data)
+          TimeHelper.time(data, timestamp_format(shape))
+        end
+
+        def apply_union(data, shape, type)
+          shape = shape_reference(shape)
+          key, value = data.flatten
+          return if key.nil?
+
+          if (member = member_with_json_name(key, shape))
+            apply_union_member(member.name, value, shape, type)
+          elsif shape.name_by_member_name?(key)
+            apply_union_member(key, value, shape, type)
+          else
+            shape.member_type(:unknown).new(key, value)
+          end
+        end
+
+        def apply_union_member(key, value, shape, type)
+          member_name = shape.name_by_member_name(key)
+          type = shape.member_type(member_name) if type.nil?
+          type.new(apply(value, shape.member(member_name)))
+        end
+
+        def extract_blob(data)
+          Base64.strict_encode64(data.is_a?(String) ? data : data.read)
+        end
+
+        def extract_list(data, shape)
+          shape = shape_reference(shape)
+          data.collect { |v| extract(v, shape.member) }
+        end
+
+        def extract_map(data, shape)
+          shape = shape_reference(shape)
+          data.each.with_object({}) { |(k, v), h| h[k.to_s] = extract(v, shape.value) }
+        end
+
+        def extract_structure(data, shape, opts)
+          shape = shape_reference(shape)
+          data.to_h.each_with_object({}) do |(k, v), o|
+            next unless shape.member?(k)
+
+            member_shape = shape.member(k)
+            member_name = resolve_member_name(member_shape, opts)
+            pp member_name
+            o[member_name] = extract(v, member_shape, opts)
+          end
+        end
+
+        def extract_timestamp(data, shape, opts)
+          return unless data.is_a?(Time)
+
+          trait = opts[:use_timestamp_format] ? timestamp_format(shape) : 'epoch-seconds'
+          TimeHelper.time(data, trait)
+        end
+
+        # rubocop:disable Metrics/AbcSize
+        def extract_union(data, shape, opts)
+          h = {}
+          shape = shape_reference(shape)
+          if data.is_a?(Schema::Union)
+            member_shape = shape.member_by_type(data.class)
+            member_name = resolve_member_name(member_shape, opts)
+            h[member_name] = extract(data, member_shape).value
+          else
+            key, value = data.first
+            if shape.member?(key)
+              member_shape = shape.member(key)
+              member_name = resolve_member_name(member_shape, opts)
+              h[member_name] = extract(value, member_shape)
+            end
+          end
+          h
+        end
+        # rubocop:enable Metrics/AbcSize
+
+        def member_name(shape, key)
+          return unless shape.name_by_member_name?(key) || shape.member?(key.to_sym)
+
+          shape.name_by_member_name(key) || key.to_sym
+        end
+
+        def member_with_json_name(name, shape)
+          shape.members.values.find do |v|
+            v.traits['smithy.api#jsonName'] == name if v.traits.include?('smithy.api#jsonName')
+          end
+        end
+
+        def resolve_member_name(member_shape, opts)
+          if opts[:use_json_name] && member_shape.traits['smithy.api#jsonName']
+            member_shape.traits['smithy.api#jsonName']
+          else
+            member_shape.name
+          end
+        end
+
+        def shape_reference(shape)
+          shape.is_a?(Shapes::MemberShape) ? shape.shape : shape
+        end
+
+        # The following steps are taken to determine the format of timestamp:
+        # Use the timestampFormat trait of the member, if present.
+        # Use the timestampFormat trait of the shape, if present.
+        # If none of the above applies, use epoch-seconds as default
+        def timestamp_format(shape)
+          if shape.traits['smithy.api#timestampFormat']
+            shape.traits['smithy.api#timestampFormat']
+          elsif shape.shape.traits['smithy.api#timestampFormat']
+            shape.shape.traits['smithy.api#timestampFormat']
+          else
+            'epoch-seconds'
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -43,7 +43,7 @@ def _to_h_array(obj)
       end
     end
 
-    # An empty Struct that includes the {Client::Structure} module.
+    # An empty Struct that includes the {Schema::Structure} module.
     EmptyStructure = Struct.new do
       include Smithy::Schema::Structure
     end