Make the router's routing state authoritative instead of sniffed

lib/cli/ui/stdout_router.rb

@@ -11,11 +11,42 @@
 module CLI
   module UI
     module StdoutRouter
+      WRITE_WITHOUT_CLI_UI = :write_without_cli_ui
+      COMPATIBILITY_WRITE_IVAR = :@__cli_ui_stdout_router_compatibility_write
+      private_constant :COMPATIBILITY_WRITE_IVAR
+
+      # Defined here rather than on the singleton class, where callers had no
+      # way to name it.
+      NotEnabled = Class.new(StandardError)
+
+      class CompatibilityWrite
+        #: Method
+        attr_reader :installed_method
+
+        #: Method
+        attr_reader :original_write
+
+        #: (Method installed_method, Method original_write) -> void
+        def initialize(installed_method, original_write)
+          @installed_method = installed_method
+          @original_write = original_write
+        end
+      end
+      private_constant :CompatibilityWrite
+
       class Writer
-        #: (io_like stream, Symbol name) -> void
-        def initialize(stream, name)
+        #: Method
+        attr_reader :original_write
+
+        # `original_write` must be the stream's pre-routing `write`; it is
+        # required because resolving it from an already-routed stream would
+        # stack a second Writer on the first. Holding the method itself keeps
+        # an in-flight write working after `deactivate`.
+        #: (io_like stream, Symbol name, Method original_write) -> void
+        def initialize(stream, name, original_write)
           @stream = stream
           @name = name
+          @original_write = original_write
         end
 
         #: (*Object args) -> Integer
@@ -38,7 +69,7 @@ def write(*args)
           end
 
           stream_args = prepend_id(@stream, strs) #: as untyped
-          ret = @stream.write_without_cli_ui(*stream_args) #: as Integer
+          ret = @original_write.call(*stream_args) #: as Integer
           if (dup = StdoutRouter.duplicate_output_to)
             begin
               dup_args = prepend_id(dup, strs) #: as untyped
@@ -207,38 +238,39 @@ def initialize(
         def run
           require 'stringio'
 
-          StdoutRouter.assert_enabled!
+          StdoutRouter.with_enabled do
+            prev_frame_inset = Thread.current[:no_cliui_frame_inset]
+            prev_hook = Thread.current[:cliui_output_hook]
+            Thread.current[:cliui_current_capture] = self
 
-          Thread.current[:cliui_current_capture] = self
-
-          prev_frame_inset = Thread.current[:no_cliui_frame_inset]
-          prev_hook = Thread.current[:cliui_output_hook]
+            begin
+              if Thread.current.respond_to?(:report_on_exception)
+                Thread.current.report_on_exception = false
+              end
 
-          if Thread.current.respond_to?(:report_on_exception)
-            Thread.current.report_on_exception = false
-          end
+              self.class.with_stdin_masked do
+                Thread.current[:no_cliui_frame_inset] = !@with_frame_inset
+                Thread.current[:cliui_output_hook] = ->(data, stream) do
+                  stream = :stdout if @merged_output
+                  case stream
+                  when :stdout
+                    @out.write(data)
+                    @duplicate_output_to.write(data)
+                  when :stderr
+                    @err.write(data)
+                  else raise
+                  end
+                  print_captured_output # suppress writing to terminal by default
+                end
 
-          self.class.with_stdin_masked do
-            Thread.current[:no_cliui_frame_inset] = !@with_frame_inset
-            Thread.current[:cliui_output_hook] = ->(data, stream) do
-              stream = :stdout if @merged_output
-              case stream
-              when :stdout
-                @out.write(data)
-                @duplicate_output_to.write(data)
-              when :stderr
-                @err.write(data)
-              else raise
+                @block.call
               end
-              print_captured_output # suppress writing to terminal by default
+            ensure
+              Thread.current[:cliui_output_hook] = prev_hook
+              Thread.current[:no_cliui_frame_inset] = prev_frame_inset
+              Thread.current[:cliui_current_capture] = nil
             end
-
-            @block.call
           end
-        ensure
-          Thread.current[:cliui_output_hook] = prev_hook
-          Thread.current[:no_cliui_frame_inset] = prev_frame_inset
-          Thread.current[:cliui_current_capture] = nil
         end
 
         #: -> String
@@ -307,10 +339,6 @@ def synchronize(&block)
       end
 
       class << self
-        WRITE_WITHOUT_CLI_UI = :write_without_cli_ui
-
-        NotEnabled = Class.new(StandardError)
-
         #: io_like?
         attr_accessor :duplicate_output_to
 
@@ -336,65 +364,176 @@ def current_id
 
         #: -> void
         def assert_enabled!
-          raise NotEnabled unless enabled?
+          raise NotEnabled unless current_streams.all? { |stream,| routed?(stream) }
         end
 
+        # Routes only streams not already routed by this module, then unroutes
+        # only those streams when the scope exits.
         #: [T] { -> T } -> T
         def with_enabled(&block)
-          enable
+          activated = activate_current_streams
           yield
         ensure
-          disable
+          activated&.reverse_each { |stream| deactivate(stream) }
         end
 
         # TODO: remove this
         #: -> void
         def ensure_activated
-          enable unless enabled?
+          enable
         end
 
+        # Routes the current $stdout/$stderr; returns whether anything changed.
+        # Additive and per-stream so a half-routed process can recover without
+        # unrouting a stream that was only replaced in the globals temporarily.
         #: -> bool
         def enable
-          return false if enabled?($stdout) || enabled?($stderr)
+          !activate_current_streams.empty?
+        end
 
-          activate($stdout, :stdout)
-          activate($stderr, :stderr)
-          true
+        #: (?io_like stream) -> bool
+        def routed?(stream = $stdout)
+          !route_writers[stream].nil?
         end
 
+        # Kept for compatibility. `routed?` names the membership check more
+        # precisely.
         #: (?io_like stream) -> bool
         def enabled?(stream = $stdout)
-          stream.respond_to?(WRITE_WITHOUT_CLI_UI)
+          routed?(stream)
         end
 
-        #: -> bool
-        def disable
-          return false unless enabled?($stdout) && enabled?($stderr)
+        # Returns a snapshot so callers can inspect routing without gaining
+        # access to the mutable registry or its Writer values.
+        #: -> Array[io_like]
+        def routed_streams
+          streams = [] #: Array[io_like]
+          route_writers.each_key do |stream|
+            streams << stream if routed?(stream)
+          end
+          streams
+        end
 
-          deactivate($stdout)
-          deactivate($stderr)
-          true
+        # Unroutes `stream`, or every routed stream when none is given; returns
+        # whether anything changed.
+        #: (?io_like? stream) -> bool
+        def disable(stream = nil)
+          routed = stream.nil? ? routed_streams : [stream].select { |candidate| routed?(candidate) }
+          routed.each { |routed_stream| deactivate(routed_stream) }
+          !routed.empty?
+        end
+
+        # Writes past the router: no frame inset, capture hooks, or output
+        # duplication. Prefer this to calling WRITE_WITHOUT_CLI_UI directly,
+        # which exists only after a stream has been routed at least once.
+        #: (io_like stream, *Object args) -> Integer
+        def write_without_routing(stream, *args)
+          compatibility_write = owned_compatibility_write(stream)
+          original_write = compatibility_write&.original_write || route_writers[stream]&.original_write
+          write_args = args #: as untyped
+          return stream.write(*write_args) unless original_write
+
+          original_write.call(*write_args) #: as Integer
         end
 
         private
 
+        # Weak keys let a caller-owned stream and its routing wrapper disappear
+        # together if the caller drops the stream without explicitly disabling
+        # it. Ruby 3.2's WeakMap has no `delete`, so deactivation tombstones the
+        # value; public readers filter those entries.
+        #: -> untyped
+        def route_writers
+          @route_writers ||= ObjectSpace::WeakMap.new
+        end
+
+        #: -> Hash[io_like, Symbol]
+        def current_streams
+          streams = {}.compare_by_identity #: Hash[io_like, Symbol]
+          streams[$stdout] = :stdout
+          streams[$stderr] = :stderr unless streams.key?($stderr)
+          streams
+        end
+
+        #: (Hash[io_like, Symbol] streams) -> Array[io_like]
+        def activate_streams(streams)
+          streams_to_activate = streams.reject { |stream,| routed?(stream) }
+          return [] unless streams_to_activate.all? { |stream,| compatibility_write_available?(stream) }
+
+          activated = [] #: Array[io_like]
+          begin
+            streams_to_activate.each do |stream, streamname|
+              activate(stream, streamname)
+              activated << stream
+            end
+            activated
+          rescue
+            activated.reverse_each { |stream| deactivate(stream) }
+            raise
+          end
+        end
+
+        #: -> Array[io_like]
+        def activate_current_streams
+          activate_streams(current_streams)
+        end
+
+        #: (io_like stream) -> bool
+        def compatibility_write_available?(stream)
+          !stream.respond_to?(WRITE_WITHOUT_CLI_UI, true) || !owned_compatibility_write(stream).nil?
+        end
+
         #: (io_like stream) -> void
         def deactivate(stream)
+          original_write = route_writers[stream]&.original_write
+          return unless original_write
+
           sc = stream.singleton_class
+
           sc.send(:remove_method, :write)
-          sc.send(:alias_method, :write, WRITE_WITHOUT_CLI_UI)
+          # Restore only a `write` we displaced; a stream that inherited its
+          # `write` gets the class method back.
+          sc.send(:define_method, :write, original_write) if original_write&.owner == sc
+          route_writers[stream] = nil
         end
 
         #: (io_like stream, Symbol streamname) -> void
         def activate(stream, streamname)
-          writer = StdoutRouter::Writer.new(stream, streamname)
-
-          raise if stream.respond_to?(WRITE_WITHOUT_CLI_UI)
+          original_write = stream.method(:write)
+          writer = StdoutRouter::Writer.new(stream, streamname, original_write)
 
-          stream.singleton_class.send(:alias_method, WRITE_WITHOUT_CLI_UI, :write)
+          install_compatibility_write(stream, original_write)
           stream.define_singleton_method(:write) do |*args|
             writer.write(*args)
           end
+          route_writers[stream] = writer
+        end
+
+        # The supported bypass is `write_without_routing`; this is the same
+        # thing under the name out-of-tree callers already reach for. Closing
+        # over the pre-routing method keeps it a bypass after disable, even if
+        # a later `write` wrapper calls the compatibility method itself.
+        #: (io_like stream, Method original_write) -> void
+        def install_compatibility_write(stream, original_write)
+          return if owned_compatibility_write(stream)
+
+          stream.singleton_class.send(:define_method, WRITE_WITHOUT_CLI_UI) do |*args|
+            write_args = args #: as untyped
+            original_write.call(*write_args)
+          end
+          installed_method = stream.method(WRITE_WITHOUT_CLI_UI)
+          compatibility_write = CompatibilityWrite.new(installed_method, original_write)
+          stream.instance_variable_set(COMPATIBILITY_WRITE_IVAR, compatibility_write)
+        end
+
+        #: (io_like stream) -> CompatibilityWrite?
+        def owned_compatibility_write(stream)
+          compatibility_write = stream.instance_variable_get(COMPATIBILITY_WRITE_IVAR)
+          return unless compatibility_write.is_a?(CompatibilityWrite)
+          return unless stream.respond_to?(WRITE_WITHOUT_CLI_UI, true)
+          return unless stream.method(WRITE_WITHOUT_CLI_UI) == compatibility_write.installed_method
+
+          compatibility_write
         end
       end
     end

test/cli/ui/printer_test.rb

@@ -6,8 +6,7 @@ module CLI
   module UI
     class PrinterTest < Minitest::Test
       def test_puts_color
-        out, _ = capture_io do
-          CLI::UI::StdoutRouter.ensure_activated
+        out, _ = capture_io_with_router do
           assert(Printer.puts('foo', frame_color: :red))
         end
 
@@ -20,8 +19,7 @@ def test_puts_color_frame
         out = nil
         capture_io do
           Frame.open('test') do
-            out, _ = capture_io do
-              CLI::UI::StdoutRouter.ensure_activated
+            out, _ = capture_io_with_router do
               Printer.puts('foo', frame_color: :red)
             end
           end
@@ -37,8 +35,7 @@ def test_frame_with_long_texts
         out = nil
         capture_io do
           Frame.open(overlong_preamble, success_text: overlong_suffix) do
-            out, _ = capture_io do
-              CLI::UI::StdoutRouter.ensure_activated
+            out, _ = capture_io_with_router do
               Printer.puts('foo', frame_color: :red)
             end
           end

test/cli/ui/progress_test.rb

@@ -33,8 +33,7 @@ def test_with_title
       end
 
       def test_updating_title
-        out, err = capture_io do
-          CLI::UI::StdoutRouter.ensure_activated
+        out, err = capture_io_with_router do
           Progress.progress do |bar|
             3.times do |i|
               bar.update_title("Title #{i}")