rbs_protobuf 0.3.0 → 1.0.0

data/lib/rbs_protobuf/translator/protobuf_gem.rb

@@ -1,13 +1,30 @@
 module RBSProtobuf
   module Translator
     class ProtobufGem < Base
+      FIELD_ARRAY = Name::Class.new(TypeName("::Protobuf::Field::FieldArray"))
+
+      FIELD_HASH = Name::Class.new(TypeName("::Protobuf::Field::FieldHash"))
+
+      ENUM = Name::Class.new(TypeName("::Protobuf::Enum"))
+
+      MESSAGE = Name::Class.new(TypeName("::Protobuf::Message"))
+
+      TO_PROTO = Name::Interface.new(TypeName("_ToProto"))
+
+      FIELD_ARRAY_a = Name::Alias.new(TypeName("::Protobuf::field_array"))
+
+      FIELD_HASH_a = Name::Alias.new(TypeName("::Protobuf::field_hash"))
+
       attr_reader :stderr
 
-      def initialize(input, upcase_enum:, nested_namespace:, extension:, stderr: STDERR)
+      attr_reader :accept_nil_writer
+
+      def initialize(input, upcase_enum:, nested_namespace:, extension:, accept_nil_writer:, stderr: STDERR)
         super(input)
         @upcase_enum = upcase_enum
         @nested_namespace = nested_namespace
         @extension = extension
+        @accept_nil_writer = accept_nil_writer
         @stderr = stderr
       end
 
@@ -90,26 +107,22 @@ module RBSProtobuf
           end
         end
 
-        file.extension.group_by(&:extendee).each.with_index do |(name, extensions), index|
+        file.extension.each.with_index do |extension, index|
           if ignore_extension?
             if print_extension_message?
-              stderr.puts "Extension for `#{name}` ignored in `#{file.name}`; Set RBS_PROTOBUF_EXTENSION env var to generate RBS for extensions."
+              stderr.puts "Extension for `#{extension.extendee}` ignored in `#{file.name}`; Set RBS_PROTOBUF_EXTENSION env var to generate RBS for extensions."
             end
           else
-            exts = extension_to_decl(name,
-                                     extensions,
-                                     prefix: RBS::Namespace.root,
-                                     source_code_info: source_code_info,
-                                     path: [7, index])
+            ext = extension_to_decl(extension, prefix: RBS::Namespace.root, source_code_info: source_code_info, path: [7, index])
 
             if print_extension?
               stderr.puts "#=========================================================="
               stderr.puts "# Printing RBS for extensions from #{file.name}"
               stderr.puts "#"
-              RBS::Writer.new(out: stderr).write(exts)
+              RBS::Writer.new(out: stderr).write([ext])
               stderr.puts
             else
-              decls.push(*exts)
+              decls.push(ext)
             end
           end
         end
@@ -119,37 +132,20 @@ module RBSProtobuf
         end.string
       end
 
-      def message_base_class
-        RBS::AST::Declarations::Class::Super.new(
-          name: RBS::TypeName.new(
-            name: :Message,
-            namespace: RBS::Namespace.parse("::Protobuf")
-          ),
-          args: [],
-          location: nil
-        )
-      end
-
-      def repeated_field_type(type, wtype = type)
-        factory.instance_type(
-          factory.type_name("::Protobuf::Field::FieldArray"),
-          type,
-          wtype
-        )
-      end
-
       def message_to_decl(message, prefix:, message_path:, source_code_info:, path:)
         class_name = ActiveSupport::Inflector.upcase_first(message.name)
 
         RBS::AST::Declarations::Class.new(
           name: RBS::TypeName.new(name: class_name.to_sym, namespace: prefix),
-          super_class: message_base_class,
+          super_class: MESSAGE.super_class,
           type_params: [],
           location: nil,
           comment: comment_for_path(source_code_info, path, options: message.options),
           members: [],
           annotations: []
         ).tap do |class_decl|
+          class_instance_type = factory.instance_type(RBS::TypeName.new(name: class_decl.name.name, namespace: RBS::Namespace.empty))
+
           maps = {}
 
           message.nested_type.each_with_index do |nested_type, index|
@@ -177,63 +173,17 @@ module RBSProtobuf
             )
           end
 
-          field_read_types = {}
-          field_write_types = {}
+          # @type var field_types: Hash[Symbol, [RBS::Types::t, Array[RBS::Types::t], RBS::Types::t]]
+          field_types = {}
 
           message.field.each_with_index do |field, index|
             field_name = field.name.to_sym
             comment = comment_for_path(source_code_info, path + [2, index], options: field.options)
 
-            read_type, write_type = field_type(field, maps)
-
-            field_read_types[field_name] = read_type
-            field_write_types[field_name] = write_type
-
-            if read_type == write_type
-              class_decl.members << RBS::AST::Members::AttrAccessor.new(
-                name: field_name,
-                type: read_type,
-                comment: comment,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
-              )
-            else
-              class_decl.members << RBS::AST::Members::AttrReader.new(
-                name: field_name,
-                type: read_type,
-                comment: comment,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
-              )
-
-              class_decl.members << RBS::AST::Members::AttrWriter.new(
-                name: field_name,
-                type: write_type,
-                comment: comment,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
-              )
-            end
+            read_type, write_types, init_type = field_type(field, maps)
+            field_types[field_name] = [read_type, write_types, init_type]
 
-            class_decl.members << RBS::AST::Members::MethodDefinition.new(
-              name: :"#{field_name}!",
-              types: [
-                factory.method_type(
-                  type: factory.function(factory.optional_type(read_type))
-                )
-              ],
-              annotations: [],
-              comment: nil,
-              location: nil,
-              overload: false,
-              kind: :instance
-            )
+            add_field(class_decl.members, name: field_name, read_type: read_type, write_types: write_types, comment: comment)
           end
 
           class_decl.members << RBS::AST::Members::MethodDefinition.new(
@@ -241,8 +191,9 @@ module RBSProtobuf
             types: [
               factory.method_type(
                 type: factory.function().update(
-                  optional_keywords: field_write_types.transform_values {|ty|
-                    factory.param(ty)
+                  optional_keywords: field_types.transform_values {|pair|
+                    _, _, init_type = pair
+                    factory.param(init_type)
                   }
                 )
               )
@@ -254,15 +205,17 @@ module RBSProtobuf
             kind: :instance
           )
 
-          unless field_read_types.empty?
+          unless field_types.empty?
             class_decl.members << RBS::AST::Members::MethodDefinition.new(
               name: :[],
               types:
-                field_read_types.keys.map do |key|
+                field_types.map do |field_name, pair|
+                  read_type, _ = pair
+
                   factory.method_type(
-                    type: factory.function(field_read_types[key]).update(
+                    type: factory.function(read_type).update(
                       required_positionals: [
-                        factory.param(factory.literal_type(key))
+                        factory.param(factory.literal_type(field_name))
                       ]
                     )
                   )
@@ -282,21 +235,34 @@ module RBSProtobuf
               overload: false,
               kind: :instance
             )
-          end
 
-          unless field_write_types.empty?
             class_decl.members << RBS::AST::Members::MethodDefinition.new(
               name: :[]=,
               types:
-                field_write_types.keys.map do |key|
-                  factory.method_type(
-                    type: factory.function(field_write_types[key]).update(
-                      required_positionals: [
-                        factory.literal_type(key),
-                        field_write_types[key]
-                      ].map {|t| factory.param(t) }
-                    )
-                  )
+                field_types.flat_map do |field_name, pair|
+                  read_type, write_types = pair
+
+                  [read_type, *write_types].map do |type|
+                    if (type_param, type_var = interface_type?(type))
+                      factory.method_type(
+                        type: factory.function(type_var).update(
+                          required_positionals: [
+                            factory.literal_type(field_name),
+                            type_var
+                          ].map {|t| factory.param(t) }
+                        )
+                      ).update(type_params: [type_param])
+                    else
+                      factory.method_type(
+                        type: factory.function(type).update(
+                          required_positionals: [
+                            factory.literal_type(field_name),
+                            type
+                          ].map {|t| factory.param(t) }
+                        )
+                      )
+                    end
+                  end
                 end +
                   [
                     factory.method_type(
@@ -333,82 +299,310 @@ module RBSProtobuf
               )
             end
           end
+
+          class_decl.members << RBS::AST::Declarations::Interface.new(
+            name: TO_PROTO.name,
+            type_params: [],
+            members: [],
+            annotations: [],
+            comment: nil,
+            location: nil
+          ).tap do |interface_decl|
+            interface_decl.members << RBS::AST::Members::MethodDefinition.new(
+              name: :to_proto,
+              types: [
+                factory.method_type(
+                  type: factory.function(class_instance_type)
+                )
+              ],
+              annotations: [],
+              comment: nil,
+              location: nil,
+              overload: false,
+              kind: :instance
+            )
+          end
+
+          class_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("init"),
+            type_params: [],
+            type: factory.union_type(class_instance_type, TO_PROTO[]),
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `#initialize` parameter.", location: nil),
+            location: nil
+          )
+
+          class_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("field_array"),
+            type_params: [],
+            type: FIELD_ARRAY[
+              class_instance_type,
+              factory.union_type(class_instance_type, TO_PROTO[])
+            ],
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `repeated` field.", location: nil),
+            location: nil
+          )
+
+          class_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("field_hash"),
+            type_params: [RBS::AST::TypeParam.new(name: :KEY, variance: :invariant, upper_bound: nil, location: nil)],
+            type: FIELD_HASH[
+              factory.type_var(:KEY),
+              class_instance_type,
+              factory.union_type(class_instance_type, TO_PROTO[])
+            ],
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `map` field.", location: nil),
+            location: nil
+          )
+
+          class_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("array"),
+            type_params: [],
+            type: RBS::BuiltinNames::Array.instance_type(factory.union_type(class_instance_type, TO_PROTO[])),
+            annotations: [],
+            comment: nil,
+            location: nil
+          )
+
+          class_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("hash"),
+            type_params: [RBS::AST::TypeParam.new(name: :KEY, variance: :invariant, upper_bound: nil, location: nil)],
+            type: RBS::BuiltinNames::Hash.instance_type(
+              factory.type_var(:KEY),
+              factory.union_type(class_instance_type, TO_PROTO[])
+            ),
+            annotations: [],
+            comment: nil,
+            location: nil
+          )
         end
       end
 
+      def message_to_proto_type(type)
+        namespace = type.name.to_namespace
+        RBS::Types::Interface.new(
+          name: RBS::TypeName.new(name: :_ToProto, namespace: namespace),
+          args: [],
+          location: nil
+        )
+      end
+
+      def message_init_type(type)
+        RBS::Types::Alias.new(
+          name: RBS::TypeName.new(name: :init, namespace: type.name.to_namespace),
+          args: [],
+          location: nil
+        )
+      end
+
+      def message_field_array_type(type)
+        RBS::Types::Alias.new(
+          name: RBS::TypeName.new(name: :field_array, namespace: type.name.to_namespace),
+          args: [],
+          location: nil
+        )
+      end
+
+      def message_array_type(type)
+        RBS::Types::Alias.new(
+          name: RBS::TypeName.new(name: :array, namespace: type.name.to_namespace),
+          args: [],
+          location: nil
+        )
+      end
+
+      def message_hash_type(type, key)
+        RBS::Types::Alias.new(
+          name: RBS::TypeName.new(name: :hash, namespace: type.name.to_namespace),
+          args: [key],
+          location: nil
+        )
+      end
+
+      def message_field_hash_type(type, key)
+        RBS::Types::Alias.new(
+          name: RBS::TypeName.new(name: :field_hash, namespace: type.name.to_namespace),
+          args: [key],
+          location: nil
+        )
+      end
+
       def field_type(field, maps)
-        case
-        when field.type == FieldDescriptorProto::Type::TYPE_MESSAGE
-          if maps.key?(field.type_name)
-            key_field, value_field = maps[field.type_name]
-
-            key_type_r, _ = field_type(key_field, maps)
-            value_type_r, value_type_w = field_type(value_field, maps)
-
-            hash_type = factory.instance_type(
-              factory.type_name("::Protobuf::Field::FieldHash"),
-              key_type_r,
-              factory.unwrap_optional(value_type_r),
-              factory.unwrap_optional(value_type_w)
-            )
+        # @type var triple: [RBS::Types::t, Array[RBS::Types::t], RBS::Types::t]
+        triple =
+          case
+          when field.type == FieldDescriptorProto::Type::TYPE_MESSAGE
+            if maps.key?(field.type_name)
+              key_field, value_field = maps[field.type_name]
+
+              key_type_r, _ = field_type(key_field, maps)
+              value_type_r, value_write_types = field_type(value_field, maps)
+
+              value_type_r = factory.unwrap_optional(value_type_r)
+              value_write_types = value_write_types.map {|type| factory.unwrap_optional(type) }
+
+              case value_field.type
+              when FieldDescriptorProto::Type::TYPE_MESSAGE, FieldDescriptorProto::Type::TYPE_ENUM
+                value_type_r.is_a?(RBS::Types::ClassInstance) or raise
+                [
+                  message_field_hash_type(value_type_r, key_type_r),
+                  [message_hash_type(value_type_r, key_type_r)],
+                  message_hash_type(value_type_r, key_type_r)
+                ]
+              else
+                hash_type = FIELD_HASH[
+                  key_type_r,
+                  value_type_r,
+                  factory.union_type(value_type_r, *value_write_types)
+                ]
+
+                [
+                  FIELD_HASH_a[key_type_r, value_type_r],
+                  [RBS::BuiltinNames::Hash.instance_type(key_type_r, value_type_r)],
+                  RBS::BuiltinNames::Hash.instance_type(key_type_r, value_type_r)
+                ]
+              end
+            else
+              type = message_type(field.type_name)
 
-            [
-              hash_type,
-              hash_type
-            ]
-          else
+              case field.label
+              when FieldDescriptorProto::Label::LABEL_OPTIONAL
+                [
+                  factory.optional_type(type),
+                  [
+                    factory.optional_type(message_to_proto_type(type))
+                  ],
+                  factory.optional_type(message_init_type(type))
+                ]
+              when FieldDescriptorProto::Label::LABEL_REPEATED
+                [
+                  message_field_array_type(type),
+                  [
+                    message_array_type(type)
+                  ],
+                  message_array_type(type)
+                ]
+              else
+                [
+                  type,
+                  [message_to_proto_type(type)],
+                  message_init_type(type)
+                ]
+              end
+            end
+          when field.type == FieldDescriptorProto::Type::TYPE_ENUM
             type = message_type(field.type_name)
-
-            case field.label
-            when FieldDescriptorProto::Label::LABEL_OPTIONAL
-              type = factory.optional_type(type)
-              [type, type]
-            when FieldDescriptorProto::Label::LABEL_REPEATED
-              type = repeated_field_type(type)
-              [type, type]
+            enum_namespace = type.name.to_namespace
+            values = factory.alias_type(RBS::TypeName.new(name: :values, namespace: enum_namespace))
+
+            if field.label == FieldDescriptorProto::Label::LABEL_REPEATED
+              [
+                message_field_array_type(type),
+                [message_array_type(type)],
+                message_array_type(type)
+              ]
+            else
+              [
+                type,
+                [values],
+                message_init_type(type)
+              ]
+            end
+          else
+            type = base_type(field.type)
+
+            if field.label == FieldDescriptorProto::Label::LABEL_REPEATED
+              [
+                FIELD_ARRAY_a[type],
+                [RBS::BuiltinNames::Array.instance_type(type)],
+                RBS::BuiltinNames::Array.instance_type(type)
+              ]
             else
-              [type, factory.optional_type(type)]
+              [type, [], type]
             end
           end
-        when field.type == FieldDescriptorProto::Type::TYPE_ENUM
-          type = message_type(field.type_name)
-          enum_namespace = type.name.to_namespace
 
-          wtype = factory.union_type(
-            type,
-            factory.alias_type(RBS::TypeName.new(name: :values, namespace: enum_namespace))
-          )
-
-          if field.label == FieldDescriptorProto::Label::LABEL_REPEATED
-            type = repeated_field_type(type, wtype)
+        if accept_nil_writer
+          read_type, write_types, init_type = triple
+          [
+            read_type,
+            ([factory.optional_type(read_type)] + write_types.map {|t| factory.optional_type(t) }).uniq,
+            factory.optional_type(init_type)
+          ]
+        else
+          triple
+        end
+      end
 
+      def interface_type?(type)
+        case
+        when type.is_a?(RBS::Types::Interface)
+          [
+            RBS::AST::TypeParam.new(name: :M, upper_bound: type, variance: :invariant, location: nil),
+            factory.type_var(:M)
+          ]
+        when type.is_a?(RBS::Types::Optional)
+          if (type = type.type).is_a?(RBS::Types::Interface)
             [
-              type,
-              type
-            ]
-          else
-            [
-              type,
-              factory.optional_type(wtype)
+              RBS::AST::TypeParam.new(name: :M, upper_bound: type, variance: :invariant, location: nil),
+              factory.optional_type(factory.type_var(:M))
             ]
           end
-        else
-          type = base_type(field.type)
-
-          if field.label == FieldDescriptorProto::Label::LABEL_REPEATED
-            type = repeated_field_type(type)
-            [type, type]
-          else
-            [type, factory.optional_type(type)]
-          end
         end
       end
 
-      def enum_base_class
-        RBS::AST::Declarations::Class::Super.new(
-          name: factory.type_name("::Protobuf::Enum"),
-          args: [],
-          location: nil
+      def add_field(members, name:, read_type:, write_types:, comment:)
+        members << RBS::AST::Members::AttrAccessor.new(
+          name: name,
+          type: read_type,
+          comment: comment,
+          location: nil,
+          annotations: [],
+          ivar_name: false,
+          kind: :instance
+        )
+
+        unless write_types.empty?
+          members << RBS::AST::Members::MethodDefinition.new(
+            name: :"#{name}=",
+            types:
+              write_types.map do |write_type|
+                if (type_param, type = interface_type?(write_type))
+                  factory.method_type(
+                    type: factory.function(type).update(
+                      required_positionals:[factory.param(type)]
+                    )
+                  ).update(type_params: [type_param])
+                else
+                  factory.method_type(
+                    type: factory.function(write_type).update(
+                      required_positionals:[factory.param(write_type)]
+                    )
+                  )
+                end
+              end,
+            annotations: [],
+            comment: comment,
+            location: nil,
+            overload: true,
+            kind: :instance
+          )
+        end
+
+        members << RBS::AST::Members::MethodDefinition.new(
+          name: :"#{name}!",
+          types: [
+            factory.method_type(
+              type: factory.function(factory.optional_type(read_type))
+            )
+          ],
+          annotations: [],
+          comment: nil,
+          location: nil,
+          overload: false,
+          kind: :instance
         )
       end
 
@@ -425,7 +619,7 @@ module RBSProtobuf
 
         RBS::AST::Declarations::Class.new(
           name: RBS::TypeName.new(name: enum_name.to_sym, namespace: prefix),
-          super_class: enum_base_class(),
+          super_class: ENUM.super_class,
           type_params: factory.module_type_params(),
           members: [],
           comment: comment_for_path(source_code_info, path, options: enum_type.options),
@@ -514,79 +708,117 @@ module RBSProtobuf
               location: nil
             )
           end
-        end
-      end
 
-      def extension_to_decl(extendee_name, extensions, prefix:, source_code_info:, path:)
-        class_name = message_type(extendee_name).name
+          enum_instance_type = factory.instance_type(RBS::TypeName.new(name: enum_name.to_sym, namespace: RBS::Namespace.empty))
+          values_type = factory.alias_type(RBS::TypeName.new(name: :values, namespace: RBS::Namespace.empty))
 
-        extensions.map do |field|
-          field_name = field.name.to_sym
+          enum_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("init"),
+            type_params: [],
+            type: factory.union_type(enum_instance_type, values_type),
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `#initialize` parameter.", location: nil),
+            location: nil
+          )
 
-          RBS::AST::Declarations::Class.new(
-            name: class_name,
-            super_class: nil,
+          enum_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("field_array"),
             type_params: [],
-            location: nil,
+            type: FIELD_ARRAY[
+              enum_instance_type,
+              factory.union_type(enum_instance_type, values_type)
+            ],
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `repeated` field.", location: nil),
+            location: nil
+          )
+
+          enum_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("field_hash"),
+            type_params: [RBS::AST::TypeParam.new(name: :KEY, variance: :invariant, upper_bound: nil, location: nil)],
+            type: FIELD_HASH[
+              factory.type_var(:KEY),
+              enum_instance_type,
+              factory.union_type(enum_instance_type, values_type)
+            ],
+            annotations: [],
+            comment: RBS::AST::Comment.new(string: "The type of `map` field.", location: nil),
+            location: nil
+          )
+
+          enum_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("array"),
+            type_params: [],
+            type: RBS::BuiltinNames::Array.instance_type(factory.union_type(enum_instance_type, values_type)),
+            annotations: [],
             comment: nil,
-            members: [],
-            annotations: []
-          ).tap do |class_decl|
-            read_type, write_type = field_type(field, {})
+            location: nil
+          )
 
-            if read_type == write_type
-              class_decl.members << RBS::AST::Members::AttrAccessor.new(
-                name: field_name,
-                type: read_type,
-                comment: nil,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
-              )
-            else
-              class_decl.members << RBS::AST::Members::AttrReader.new(
-                name: field_name,
-                type: read_type,
-                comment: nil,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
-              )
+          enum_decl.members << RBS::AST::Declarations::Alias.new(
+            name: TypeName("hash"),
+            type_params: [RBS::AST::TypeParam.new(name: :KEY, variance: :invariant, upper_bound: nil, location: nil)],
+            type: RBS::BuiltinNames::Hash.instance_type(
+              factory.type_var(:KEY),
+              factory.union_type(enum_instance_type, values_type)
+            ),
+            annotations: [],
+            comment: nil,
+            location: nil
+          )
+        end
+      end
 
-              class_decl.members << RBS::AST::Members::AttrWriter.new(
-                name: field_name,
-                type: write_type,
-                comment: nil,
-                location: nil,
-                annotations: [],
-                ivar_name: false,
-                kind: :instance
+      def extension_to_decl(extension, prefix:, source_code_info:, path:)
+        class_name = message_type(extension.extendee).name
+
+        comment = comment_for_path(source_code_info, path, options: extension.options)
+        field_name = extension.name.to_sym
+
+        RBS::AST::Declarations::Class.new(
+          name: class_name,
+          super_class: nil,
+          type_params: [],
+          location: nil,
+          comment: nil,
+          members: [],
+          annotations: []
+        ).tap do |class_decl|
+          read_type, write_types, _ = field_type(extension, {})
+
+          add_field(class_decl.members, name: field_name, read_type: read_type, write_types: write_types, comment: comment)
+
+          class_decl.members << RBS::AST::Members::MethodDefinition.new(
+            name: :[],
+            types: [
+              factory.method_type(
+                type: factory.function(read_type).update(
+                  required_positionals: [
+                    factory.param(factory.literal_type(field_name))
+                  ]
+                )
               )
-            end
+            ],
+            annotations: [],
+            comment: nil,
+            location: nil,
+            overload: true,
+            kind: :instance
+          )
 
-            class_decl.members << RBS::AST::Members::MethodDefinition.new(
-              name: :[],
-              types: [
+          class_decl.members << RBS::AST::Members::MethodDefinition.new(
+            name: :[]=,
+            types: [read_type, *write_types].map do |write_type|
+              if (type_param, type_var = interface_type?(write_type))
                 factory.method_type(
-                  type: factory.function(read_type).update(
+                  type: factory.function(type_var).update(
                     required_positionals: [
-                      factory.param(factory.literal_type(field_name))
+                      factory.param(factory.literal_type(field_name)),
+                      factory.param(type_var)
                     ]
                   )
-                )
-              ],
-              annotations: [],
-              comment: nil,
-              location: nil,
-              overload: true,
-              kind: :instance
-            )
-
-            class_decl.members << RBS::AST::Members::MethodDefinition.new(
-              name: :[]=,
-              types: [
+                ).update(type_params: [type_param])
+              else
                 factory.method_type(
                   type: factory.function(write_type).update(
                     required_positionals: [
@@ -595,14 +827,14 @@ module RBSProtobuf
                     ]
                   )
                 )
-              ],
-              annotations: [],
-              comment: nil,
-              location: nil,
-              overload: true,
-              kind: :instance
-            )
-          end
+              end
+            end,
+            annotations: [],
+            comment: nil,
+            location: nil,
+            overload: true,
+            kind: :instance
+          )
         end
       end