lib/tapioca/dsl/compiler.rb in tapioca-0.16.2 vs lib/tapioca/dsl/compiler.rb in tapioca-0.16.3

@@ -23,10 +23,12 @@
       attr_reader :root
 
       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
 
         sig { params(constant: Module).returns(T::Boolean) }
         def handles?(constant)
@@ -42,23 +44,54 @@
             T::Set[Module].new.compare_by_identity.merge(gather_constants),
             T.nilable(T::Set[Module]),
           )
         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
       end