typing for migration

lib/has_aux_table/migration_extensions.rb

@@ -1,23 +1,59 @@
-# typed: true
+# typed: strict
 # frozen_string_literal: true
 
 module HasAuxTable
-  class BaseTableDefinition
+  TableDefinition =
-    delegate_missing_to :inner
+    T.type_alias do
-
+      T.any(
-    def initialize(schema, base_table_name, inner)
+        ActiveRecord::ConnectionAdapters::TableDefinition,
-      @schema = schema
+        ActiveRecord::ConnectionAdapters::Table
-      @base_table_name = base_table_name
+      )
-      @inner = inner
     end
 
+  TableDefinitionWithAux =
+    T.type_alias { T.all(TableDefinition, AuxTableDefinition) }
+
+  module AuxTableDefinition
+    extend T::Sig
+    extend T::Helpers
+    abstract!
+
+    sig do
+      params(
+        aux_name: Symbol,
+        options: T::Hash[Symbol, T.untyped],
+        block: T.untyped
+      ).void
+    end
     def create_aux(aux_name, **options, &block)
-      @schema.create_aux_table(@base_table_name, aux_name, **options, &block)
+      schema.create_aux_table(aux_base_table_name, aux_name, **options, &block)
     end
 
-    private
+    protected
 
-    attr_reader :schema, :inner
+    sig { abstract.returns(ActiveRecord::Schema) }
+    def schema
+    end
+
+    sig { abstract.returns(Symbol) }
+    def aux_base_table_name
+    end
+
+    sig do
+      params(
+        table_definition: TableDefinition,
+        schema: HasAuxTable::MigrationExtensions,
+        aux_base_table_name: Symbol
+      ).returns(TableDefinitionWithAux)
+    end
+    def self.augment(table_definition, schema, aux_base_table_name)
+      table_definition.instance_eval do
+        extend AuxTableDefinition
+        define_singleton_method(:schema) { schema }
+        define_singleton_method(:aux_base_table_name) { aux_base_table_name }
+      end
+      T.cast(table_definition, TableDefinitionWithAux)
+    end
   end
 
   module MigrationExtensions
@@ -25,21 +61,45 @@ module HasAuxTable
     extend T::Helpers
     requires_ancestor { ActiveRecord::Migration }
 
-    def create_base_table(name, type: :string, **options)
+    sig do
+      params(
+        name: Symbol,
+        type: Symbol,
+        options: T::Hash[Symbol, T.untyped],
+        block: T.proc.params(table_definition: TableDefinitionWithAux).void
+      ).void
+    end
+    def create_base_table(name, type: :string, **options, &block)
       create_table(name, **options) do |t|
         t.column :type, type, null: false
-        t = BaseTableDefinition.new(self, name, t)
+        t = AuxTableDefinition.augment(t, self, name)
-        yield t
+        block.call(t)
       end
     end
 
-    def change_base_table(name, **options)
+    sig do
+      params(
+        name: Symbol,
+        options: T::Hash[Symbol, T.untyped],
+        block: T.proc.params(table_definition: TableDefinitionWithAux).void
+      ).void
+    end
+    def change_base_table(name, **options, &block)
       change_table(name, **options) do |t|
-        yield BaseTableDefinition.new(self, name, t)
+        t = AuxTableDefinition.augment(t, self, name)
+        block.call(t)
       end
     end
 
-    def create_aux_table(base_table, name, **options)
+    sig do
+      params(
+        base_table: Symbol,
+        name: Symbol,
+        options: T::Hash[Symbol, T.untyped],
+        block: T.proc.params(table_definition: TableDefinitionWithAux).void
+      ).void
+    end
+    def create_aux_table(base_table, name, **options, &block)
       aux_table_name = "#{base_table}_#{name}_aux"
       create_table(aux_table_name, id: false, **options) do |t|
         t.references :base_table,
@@ -49,7 +109,8 @@ module HasAuxTable
                        to_table: base_table,
                        validate: true
                      }
-        yield t
+        t = AuxTableDefinition.augment(t, self, base_table)
+        block.call(t)
       end
     end
   end