Iterate over requested_constants in DSL helpers

Shopify · Aug 27, 2024 · 547b5ac · 547b5ac
1 parent 2634126
commit 547b5ac
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 2 deletions.
diff --git a/lib/tapioca/dsl/compiler.rb b/lib/tapioca/dsl/compiler.rb
@@ -25,6 +25,8 @@ class Compiler
       sig { returns(T::Hash[String, T.untyped]) }
       attr_reader :options
 
+      @@requested_constants = T.let([], T::Array[Module]) # rubocop:disable Style/ClassVars
+
       class << self
         extend T::Sig
 
@@ -44,20 +46,51 @@ def processable_constants
           )
         end
 
+        sig { params(constants: T::Array[Module]).void }
+        def requested_constants=(constants)
+          @@requested_constants = constants # rubocop:disable Style/ClassVars
+        end
+
         private
 
+        sig do
+          type_parameters(:U)
+            .params(klass: T.all(T::Class[T.anything], T.type_parameter(:U)))
+            .returns(T::Array[T.type_parameter(:U)])
+        end
+        def descendants_of(klass)
+          if @@requested_constants.any?
+            T.cast(
+              @@requested_constants.select do |k|
+                k < klass && !k.singleton_class?
+              end,
+              T::Array[T.type_parameter(:U)],
+            )
+          else
+            super
+          end
+        end
+
         sig { returns(T::Enumerable[T::Class[T.anything]]) }
         def all_classes
           @all_classes ||= T.let(
-            ObjectSpace.each_object(Class),
+            if @@requested_constants.any?
+              @@requested_constants.grep(Class)
+            else
+              ObjectSpace.each_object(Class)
+            end,
             T.nilable(T::Enumerable[T::Class[T.anything]]),
           )
         end
 
         sig { returns(T::Enumerable[Module]) }
         def all_modules
           @all_modules ||= T.let(
-            ObjectSpace.each_object(Module),
+            if @@requested_constants.any?
+              @@requested_constants.select { |k| k.is_a?(Module) }
+            else
+              ObjectSpace.each_object(Module)
+            end,
             T.nilable(T::Enumerable[Module]),
           )
         end

diff --git a/lib/tapioca/dsl/pipeline.rb b/lib/tapioca/dsl/pipeline.rb
@@ -145,6 +145,7 @@ def gather_active_compilers(requested_compilers, excluded_compilers)
         ).returns(T::Set[Module])
       end
       def gather_constants(requested_constants, requested_paths, skipped_constants)
+        Compiler.requested_constants = requested_constants
         constants = Set.new.compare_by_identity
         active_compilers.each do |compiler|
           constants.merge(compiler.processable_constants)