utils/progress.h

Go to the documentation of this file.

#ifndef UTILS_PROGRESS_H
#define UTILS_PROGRESS_H

#include "utils/env.h"
#include "utils/logger.h"
#include "utils/stringutils.h"

#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <unistd.h>
#include <sys/ioctl.h>

namespace utils
{

// TODO replace this with constexpr in C++11
#define ROTATION_IND "-\\|/"

class Progress
{
private:
        enum OutputType
        {
                DISABLED,
                TTY
                // TODO Support file output
        };

private:
        std::ostream *m_output;

        OutputType m_type;

        unsigned long m_total;

        unsigned long m_current;

        unsigned long m_barSize;

        unsigned char m_rotPosition;

        std::ofstream m_tty;

public:
        Progress(unsigned long total = 100)
                : m_current(0), m_barSize(80) /* Default size */, m_rotPosition(0)
        {
                std::string envOutput = Env::get<std::string>("UTILS_PROGRESS_OUTPUT", "STDERR");

                StringUtils::toUpper(envOutput);

                if (envOutput == "STDOUT") {
                        m_output = &std::cout;
                        m_type = TTY;

                        setSize(isatty(fileno(stdout)));
                } else if (envOutput == "STDERR") {
                        m_output = &std::cerr;
                        m_type = TTY;

                        setSize(isatty(fileno(stderr)));
                } else if (envOutput == "TTY") {
                        m_tty.open("/dev/tty"); // try unix
                        if (!m_tty)
                                m_tty.open("CON:"); // try windows

                        if (m_tty) {
                                m_output = &m_tty;
                                m_type = TTY;
                                setSize();
                        } else {
                                logWarning() << "Could not open terminal. Disabling progress bar.";
                                m_type = DISABLED;
                        }
                } else {
                        m_type = DISABLED;
                }
        }

        void setTotal(unsigned long total)
        {
                m_total = total;
        }

        void set(unsigned long current)
        {
                m_current = std::min(current, m_total);
        }

        void update(unsigned long current)
        {
                set(current);

                if (m_type == DISABLED)
                        return;

                // Calculuate the ratio of complete-to-incomplete.
                const float ratio = m_current/static_cast<float>(m_total);

                // Show the percentage complete
                (*m_output) << std::setw(3) << (int)(ratio*100) << "% [";

                // real width (without additional chars)
                const unsigned long realSize = m_barSize - 9;

                const unsigned long comChars = realSize * ratio;

                // Show the load bar
                for (unsigned int i = 0; i < comChars; i++)
                        (*m_output) << '=';

                for (unsigned int i = comChars; i < realSize; i++)
                        (*m_output) << ' ';

                (*m_output) << "] ";

                // Print rotation indicator
                (*m_output) << ROTATION_IND[m_rotPosition];
                m_rotPosition = (m_rotPosition+1) % 4;

                // go to the beginning of the line
                (*m_output) << '\r' << std::flush;
        }

        void update()
        {
                update(m_current);
        }

        void increment()
        {
                update(m_current + 1);
        }

        void clear()
        {
                if (m_type == DISABLED)
                        return;

                for (unsigned int i = 0; i < m_barSize; i++)
                        (*m_output) << ' ';
                (*m_output) << '\r' << std::flush;
        }

private:
        void setSize(bool automatic = true)
        {
                // Check if size is set in env
                unsigned long size = Env::get<unsigned long>("UTILS_PROGRESS_SIZE", 0);

                if (size > 0) {
                        m_barSize = size;
                        return;
                }

                if (!automatic)
                        // No automatic size detection (e.g. for stdout)
                        return;

                // Try to get terminal size
#ifdef TIOCGSIZE
                struct ttysize ts;
                ioctl(STDIN_FILENO, TIOCGSIZE, &ts);
                m_barSize = ts.ts_cols;
#elif defined(TIOCGWINSZ)
                struct winsize ts;
                ioctl(STDIN_FILENO, TIOCGWINSZ, &ts);
                m_barSize = ts.ws_col;
#else
                logWarning() << "Could not get terminal size, using default";
#endif
        }
};

}

#endif // UTILS_PROGRESS_H