range.hpp

/****************************************************************************
 *  hipipe library
 *  Copyright (c) 2017, Cognexa Solutions s.r.o.
 *  Copyright (c) 2018, Iterait a.s.
 *  Author(s) Filip Matzner
 *
 *  This file is distributed under the MIT License.
 *  See the accompanying file LICENSE.txt for the complete license agreement.
 ****************************************************************************/
 
#pragma once
#include <hipipe/build_config.hpp>
#ifdef HIPIPE_BUILD_PYTHON
 
#include <hipipe/core/python/utility/pyboost_is_registered.hpp>
 
#include <boost/python.hpp>
#include <range/v3/core.hpp>
 
#include <memory>
#include <stdexcept>
#include <string>
#include <typeinfo>
#include <vector>
 
namespace hipipe::python {
 
struct stop_iteration_exception : public std::runtime_error {
    stop_iteration_exception()
      : std::runtime_error{"stop iteration"}
    { }
};
 
inline void stop_iteration_translator(const stop_iteration_exception& x)
{
    PyErr_SetNone(PyExc_StopIteration);
}
 
template<typename Rng>
class range {
private:
    std::shared_ptr<Rng> rng_ptr_;
 
    // register __len__ function if it is supported
    CPP_template(int dummy = 0)(requires ranges::sized_range<const Rng>)
    static void register_len(boost::python::class_<range<Rng>>& cls)
    {
        cls.def("__len__", &range<Rng>::len<>);
    }
    CPP_template(int dummy = 0)(requires !ranges::sized_range<const Rng>)
    static void register_len(boost::python::class_<range<Rng>>&)
    {
    }
 
    // register __getitem__ function if it is supported
    CPP_template(int dummy = 0)(requires ranges::random_access_range<const Rng>)
    static void register_getitem(boost::python::class_<range<Rng>>& cls)
    {
        cls.def("__getitem__", &range<Rng>::getitem<>);
    }
    CPP_template(int dummy = 0)(requires !ranges::random_access_range<const Rng>)
    static void register_getitem(boost::python::class_<range<Rng>>&)
    {
    }
 
    // function to register the type of this class in boost::python
    // makes sure the type is registered only once
    static void register_to_python()
    {
        namespace py = boost::python;
 
        if (!utility::is_registered<range<Rng>>()) {
            std::string this_name = std::string("hipipe_") + typeid(range<Rng>).name();
            py::class_<range<Rng>> cls{this_name.c_str(), py::no_init};
            cls.def("__iter__", &range<Rng>::iter);
            register_len(cls);
            register_getitem(cls);
 
            py::class_<range<Rng>::iterator>{(this_name + "_iterator").c_str(), py::no_init}
              .def("__next__", &range<Rng>::iterator::next)
              .def("__iter__", &range<Rng>::iterator::iter);
        };
    }
 
public:
    class iterator {
    private:
        std::shared_ptr<Rng> rng_ptr_;
        ranges::iterator_t<Rng> position_;
        bool first_iteration_;
 
    public:
        iterator() = default;
 
        // Construct iterator from a range.
        explicit iterator(range& rng)
          : rng_ptr_{rng.rng_ptr_},
            position_{ranges::begin(*rng_ptr_)},
            first_iteration_{true}
        {
        }
 
        // Return a copy of this iterator.
        iterator iter()
        {
            return *this;
        }
 
        // Return the next element in the range.
        //
        // \throws stop_iteration_exception if there are no more elements.
        auto next()
        {
            // do not increment the iterator in the first iteration, just return *begin()
            if (!first_iteration_ && position_ != ranges::end(*rng_ptr_)) ++position_;
            first_iteration_ = false;
            if (position_ == ranges::end(*rng_ptr_)) throw stop_iteration_exception();
            return *position_;
        }
    };
 
    range() = default;
 
    explicit range(Rng rng)
      : rng_ptr_{std::make_shared<Rng>(std::move(rng))}
    {
        register_to_python();
    }
 
    range(std::shared_ptr<Rng> rng_ptr)
      : rng_ptr_{std::move(rng_ptr)}
    {
        register_to_python();
    }
 
    // Get python iterator.
    iterator iter()
    {
        return iterator{*this};
    }
 
    // Get an item or a slice.
    //
    // Note that when slicing, the data get copied.
    CPP_template(int dummy = 0)(requires ranges::sized_range<const Rng>)
    boost::python::object getitem(PyObject* idx_py) const
    {
        // handle slices
        if (PySlice_Check(idx_py)) {
            PySliceObject* slice = static_cast<PySliceObject*>(static_cast<void*>(idx_py));
            if (slice->step != Py_None) {
                throw std::logic_error("HiPipe python range does not support slice steps.");
            }
 
            auto handle_index = [this](PyObject* idx_py, long def_val) {
                long idx = def_val;
                if (idx_py != Py_None) {
                    idx = boost::python::extract<long>(idx_py);
                    // reverse negative index
                    if (idx < 0) idx += this->len();
                    // if it is still negative, clip
                    if (idx < 0) idx = 0;
                    // handle index larger then len
                    if (idx > this->len()) idx = len();
                }
                return idx;
            };
            long start = handle_index(slice->start, 0);
            long stop = handle_index(slice->stop, len());
            if (start > stop) start = stop;
 
            using slice_data_type = std::vector<ranges::range_value_t<Rng>>;
            slice_data_type slice_data{rng_ptr_->begin() + start, rng_ptr_->begin() + stop};
            return boost::python::object{range<slice_data_type>{std::move(slice_data)}};
        }
 
        // handle indices
        long idx = boost::python::extract<long>(idx_py);
        if (idx < 0) idx += len();
        return boost::python::object{ranges::at(*rng_ptr_, idx)};
    }
 
    // Get the size of the range.
    CPP_template(int dummy = 0)(requires ranges::sized_range<const Rng>)
    long len() const
    {
        return ranges::size(*rng_ptr_);
    }
 
};  // class range
 
}  // end namespace hipipe::python
 
#endif  // HIPIPE_BUILD_PYTHON