IOIntercept.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
#include <iostream>
#include <string>
#include <framework/logging/LogConfig.h>
 
namespace Belle2 {
  namespace IOIntercept {
    class StreamInterceptor {
    public:
      StreamInterceptor(std::ostream& stream, FILE* fileObject);
      ~StreamInterceptor();

      bool start()
      {
        // only start if we are not already running
        if (!m_capturing) m_capturing = replaceFD(m_replacementFD);
        return m_capturing;
      }

      bool finish()
      {
        if (!m_capturing) return false;
        m_capturing = false;
        return replaceFD(m_savedFD);
      }
    protected:
      void setReplacementFD(int fd) { m_replacementFD = fd; }
      static void readFD(int fd, std::string& out);
      bool replaceFD(int fileDescriptor);
      std::ostream& m_stream;
      FILE* m_fileObject;
      int m_savedFD{ -1};
      int m_replacementFD{ -1};
      bool m_capturing{false};
    };

    class KeepStream {
    public:
      KeepStream(std::ostream&, FILE*) {}
      bool start() const { return true; }
      bool finish() const { return true; }
    };

    class DiscardStream: public StreamInterceptor {
    public:
      DiscardStream(std::ostream& stream, FILE* fileObject);
    };

    class CaptureStream: public StreamInterceptor {
    public:
      CaptureStream(std::ostream& stream, FILE* fileObject);
      ~CaptureStream();
      const std::string& getOutput() const { return m_outputStr; }
      bool start();
      bool finish();
    private:
      int m_pipeReadFD{ -1};
      std::string m_outputStr;
      friend class CaptureStreamAbortHandler;
    };
 

    template<class STDOUT, class STDERR>
    class InterceptOutput {
    public:
      bool start() { return m_stdout.start() && m_stderr.start(); }
      bool finish() { return m_stdout.finish() && m_stderr.finish(); }
      const std::string& getStdOut() const
      {
        static_assert(std::is_same<STDOUT, CaptureStream>::value, "Only valid if stdout is captured using CaptureStream");
        return m_stdout.getOutput();
      }

      const std::string& getStdErr() const
      {
        static_assert(std::is_same<STDERR, CaptureStream>::value, "Only valid if stderr is captured using CaptureStream");
        return m_stderr.getOutput();
      }
 
    private:
      STDOUT m_stdout{std::cout, stdout};
      STDERR m_stderr{std::cerr, stderr};
    };

    using CaptureStdOutStdErr = InterceptOutput<CaptureStream, CaptureStream>;
    using CaptureStdOut = InterceptOutput<CaptureStream, KeepStream>;
    using CaptureStdErr = InterceptOutput<KeepStream, CaptureStream>;
    using DiscardStdOutStdErr = InterceptOutput<DiscardStream, DiscardStream>;
    using DiscardStdOut = InterceptOutput<DiscardStream, KeepStream>;
    using DiscardStdErr = InterceptOutput<KeepStream, DiscardStream>;
    using CaptureStdOutDiscardStdErr = InterceptOutput<CaptureStream, DiscardStream>;
    using DiscardStdOutCaptureStdErr = InterceptOutput<DiscardStream, CaptureStream>;

    class OutputToLogMessages: public CaptureStdOutStdErr {
    public:
      OutputToLogMessages(const std::string& name, LogConfig::ELogLevel stdoutLevel, LogConfig::ELogLevel stderrLevel,
                          int stdoutDebugLevel, int stderrDebugLevel):
        m_name(name), m_indent(name + ": "), m_stdoutLevel(stdoutLevel), m_stderrLevel(stderrLevel),
        m_stdoutDebugLevel(stdoutDebugLevel), m_stderrDebugLevel(stderrDebugLevel)
      {}

      OutputToLogMessages(const std::string& name, LogConfig::ELogLevel stdoutLevel, LogConfig::ELogLevel stderrLevel):
        OutputToLogMessages(name, stdoutLevel, stderrLevel, 100, 100)
      {}

      explicit OutputToLogMessages(const std::string& name): OutputToLogMessages(name, LogConfig::c_Info, LogConfig::c_Error)
      {}

      ~OutputToLogMessages() { finish(); }
      void setIndent(const std::string& indent) { m_indent = indent; }
      bool finish();
    private:
      const std::string m_name;
      std::string m_indent;
      LogConfig::ELogLevel m_stdoutLevel;
      LogConfig::ELogLevel m_stderrLevel;
      int m_stdoutDebugLevel;
      int m_stderrDebugLevel;
    };

    template<class T> class InterceptorScopeGuard {
    public:
      explicit InterceptorScopeGuard(T& interceptor): m_interceptor(&interceptor)
      {
        m_interceptor->start();
      }

      InterceptorScopeGuard(InterceptorScopeGuard<T>&& b): m_interceptor(b.m_interceptor)
      {
        b.m_interceptor = nullptr;
      }

      InterceptorScopeGuard(const InterceptorScopeGuard<T>&) = delete;
      InterceptorScopeGuard& operator=(const InterceptorScopeGuard&) = delete;
      ~InterceptorScopeGuard()
      {
        if (m_interceptor) m_interceptor->finish();
      }
    private:
      T* m_interceptor;
    };

    template<class T> InterceptorScopeGuard<T> start_intercept(T& interceptor)
    {
      return InterceptorScopeGuard<T> {interceptor};
    }
  }

}