Infer types from the first record

Shopify · Apr 14, 2022 · 1e96598 · 1e96598
1 parent 7872cf2
commit 1e96598
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 39 deletions.
diff --git a/lib/tapioca/dsl/compilers/frozen_record.rb b/lib/tapioca/dsl/compilers/frozen_record.rb
@@ -72,22 +72,18 @@ def decorate
           attributes = constant.attributes
           return if attributes.empty?
 
+          instance = constant.first
+
           root.create_path(constant) do |record|
             module_name = "FrozenRecordAttributeMethods"
 
             record.create_module(module_name) do |mod|
-              extra_methods = constant.instance_methods(false) - attributes.to_a.map(&:to_sym)
               attributes.each do |attribute|
-                return_type = compile_method_return_type_to_rbi(constant.instance_method(attribute))
+                return_type = instance.attributes[attribute].class.name
+                return_type = "T::Boolean" if ["FalseClass", "TrueClass"].include?(return_type)
                 mod.create_method("#{attribute}?", return_type: "T::Boolean")
                 mod.create_method(attribute.to_s, return_type: return_type)
               end
-              extra_methods.each do |method|
-                method_def = constant.instance_method(method)
-                parameters = compile_method_parameters_to_rbi(method_def)
-                return_type = compile_method_return_type_to_rbi(method_def)
-                mod.create_method(method.to_s, return_type: return_type, parameters: parameters)
-              end
             end
 
             record.create_include(module_name)

diff --git a/spec/tapioca/dsl/compilers/frozen_record_spec.rb b/spec/tapioca/dsl/compilers/frozen_record_spec.rb
@@ -73,19 +73,19 @@ class Student
                   include FrozenRecordAttributeMethods
 
                   module FrozenRecordAttributeMethods
-                    sig { returns(T.untyped) }
+                    sig { returns(String) }
                     def first_name; end
 
                     sig { returns(T::Boolean) }
                     def first_name?; end
 
-                    sig { returns(T.untyped) }
+                    sig { returns(Integer) }
                     def id; end
 
                     sig { returns(T::Boolean) }
                     def id?; end
 
-                    sig { returns(T.untyped) }
+                    sig { returns(String) }
                     def last_name; end
 
                     sig { returns(T::Boolean) }
@@ -106,25 +106,7 @@ class Student < FrozenRecord::Base
 
                   self.base_path = __dir__
 
-                  sig { returns(String) }
-                  def first_name
-                    super
-                  end
-
-                  sig { returns(String) }
-                  def last_name
-                    super
-                  end
-
-                  sig { returns(String) }
-                  def location
-                    super
-                  end
-
-                  sig { returns(Integer) }
-                  def age
-                    return super + 5
-                  end
+                  self.default_attributes = { shirt_size: :large }
 
                   sig { params(grain: Symbol).returns(String) }
                   def area(grain:)
@@ -149,11 +131,27 @@ def area(grain:)
                   last_name: Smith
                   age: 19
                   location: Ottawa, Ontario, Canada
+                  is_cool_person: no
+                  birth_date: 1867-07-01
+                  updated_at: 2014-02-24T19:08:06-05:00
+                  favourite_foods:
+                    - Pizza
+                  skills:
+                    backend: Ruby
+                    frontend: HTML
                 - id: 2
                   first_name: Dan
                   last_name:  Lord
                   age: 20
                   location: Toronto, Ontario, Canada
+                  is_cool_person: yes
+                  birth_date: 1967-07-01
+                  updated_at: 2015-02-24T19:08:06-05:00
+                  favourite_foods:
+                    - Tacos
+                  skills:
+                    backend: Ruby
+                    frontend: CSS
               YAML
 
               expected = <<~RBI
@@ -163,38 +161,71 @@ class Student
                   include FrozenRecordAttributeMethods
 
                   module FrozenRecordAttributeMethods
-                    sig { returns(::Integer) }
+                    sig { returns(Integer) }
                     def age; end
 
                     sig { returns(T::Boolean) }
                     def age?; end
 
-                    sig { params(grain: ::Symbol).returns(::String) }
-                    def area(grain:); end
+                    sig { returns(Date) }
+                    def birth_date; end
 
-                    sig { returns(::String) }
+                    sig { returns(T::Boolean) }
+                    def birth_date?; end
+
+                    sig { returns(Array) }
+                    def favourite_foods; end
+
+                    sig { returns(T::Boolean) }
+                    def favourite_foods?; end
+
+                    sig { returns(String) }
                     def first_name; end
 
                     sig { returns(T::Boolean) }
                     def first_name?; end
 
-                    sig { returns(T.untyped) }
+                    sig { returns(Integer) }
                     def id; end
 
                     sig { returns(T::Boolean) }
                     def id?; end
 
-                    sig { returns(::String) }
+                    sig { returns(T::Boolean) }
+                    def is_cool_person; end
+
+                    sig { returns(T::Boolean) }
+                    def is_cool_person?; end
+
+                    sig { returns(String) }
                     def last_name; end
 
                     sig { returns(T::Boolean) }
                     def last_name?; end
 
-                    sig { returns(::String) }
+                    sig { returns(String) }
                     def location; end
 
                     sig { returns(T::Boolean) }
                     def location?; end
+
+                    sig { returns(Symbol) }
+                    def shirt_size; end
+
+                    sig { returns(T::Boolean) }
+                    def shirt_size?; end
+
+                    sig { returns(Hash) }
+                    def skills; end
+
+                    sig { returns(T::Boolean) }
+                    def skills?; end
+
+                    sig { returns(Time) }
+                    def updated_at; end
+
+                    sig { returns(T::Boolean) }
+                    def updated_at?; end
                   end
                 end
               RBI
@@ -226,13 +257,13 @@ class Student
                   extend GeneratedRelationMethods
 
                   module FrozenRecordAttributeMethods
-                    sig { returns(T.untyped) }
+                    sig { returns(String) }
                     def course; end
 
                     sig { returns(T::Boolean) }
                     def course?; end
 
-                    sig { returns(T.untyped) }
+                    sig { returns(Integer) }
                     def id; end
 
                     sig { returns(T::Boolean) }