_schema.cc

Go to the documentation of this file.

/*
 * This file is part of afw.
 *
 * Developed for the LSST Data Management System.
 * This product includes software developed by the LSST Project
 * (https://www.lsst.org).
 * See the COPYRIGHT file at the top-level directory of this distribution
 * for details of code ownership.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"
 
#include <sstream>
 
#include "ndarray/pybind11.h"
 
#include "lsst/utils/python.h"
 
#include "lsst/afw/fits.h"
#include "lsst/afw/table/detail/Access.h"
#include "lsst/afw/table/Schema.h"
#include "lsst/afw/table/BaseRecord.h"
#include "lsst/afw/table/SchemaMapper.h"
 
namespace py = pybind11;
using namespace pybind11::literals;
 
namespace lsst {
namespace afw {
namespace table {
 
using utils::python::WrapperCollection;
 
namespace {
 
using PySchema = py::class_<Schema>;
 
using PySubSchema = py::class_<SubSchema>;
 
template <typename T>
using PyFieldBase = py::class_<FieldBase<T>>;
 
template <typename T>
using PyKeyBase = py::class_<KeyBase<T>>;
 
template <typename T>
using PyField = py::class_<Field<T>, FieldBase<T>>;
 
template <typename T>
using PyKey = py::class_<Key<T>, KeyBase<T>, FieldBase<T>>;
 
template <typename T>
using PySchemaItem = py::class_<SchemaItem<T>>;
 
// Specializations for FieldBase
 
template <typename T>
void declareFieldBaseSpecializations(PyFieldBase<T> &cls) {
    cls.def(py::init<>());
}
 
template <typename T>
void declareFieldBaseSpecializations(PyFieldBase<Array<T>> &cls) {
    cls.def(py::init<int>(), "size"_a = 0);
    cls.def("getSize", &FieldBase<Array<T>>::getSize);
    cls.def("isVariableLength", &FieldBase<Array<T>>::isVariableLength);
}
 
void declareFieldBaseSpecializations(PyFieldBase<std::string> &cls) {
    cls.def(py::init<int>(), "size"_a = -1);
    cls.def("getSize", &FieldBase<std::string>::getSize);
}
 
// Specializations for Field
 
template <typename T>
void declareFieldSpecializations(PyField<T> &cls) {
    cls.def(py::pickle(
            [](Field<T> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits());
            },
            [](py::tuple t) {
                int const NPARAMS = 3;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return Field<T>(t[0].cast<std::string>(), t[1].cast<std::string>(), t[2].cast<std::string>());
            }));
}
 
// Field<Array<T>> and Field<std::string> have the same pickle implementation
template <typename T>
void _sequenceFieldSpecializations(PyField<T> &cls) {
    cls.def(py::pickle(
            [](Field<T> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getName(), self.getDoc(), self.getUnits(), self.getSize());
            },
            [](py::tuple t) {
                int const NPARAMS = 4;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return Field<T>(t[0].cast<std::string>(), t[1].cast<std::string>(), t[2].cast<std::string>(),
                                t[3].cast<int>());
            }));
}
 
template <typename T>
void declareFieldSpecializations(PyField<Array<T>> &cls) {
    _sequenceFieldSpecializations(cls);
}
 
void declareFieldSpecializations(PyField<std::string> &cls) { _sequenceFieldSpecializations(cls); }
 
// Specializations for KeyBase
 
template <typename T>
void declareKeyBaseSpecializations(PyKeyBase<T> &) {}
 
template <typename T>
void declareKeyBaseSpecializations(PyKeyBase<Array<T>> &cls) {
    cls.def("__getitem__", [](Key<Array<T>> const &self, py::object const &index) -> py::object {
        if (py::isinstance<py::slice>(index)) {
            py::slice slice(index);
            py::size_t start = 0, stop = 0, step = 0, length = 0;
            bool valid = slice.compute(self.getSize(), &start, &stop, &step, &length);
            if (!valid) throw py::error_already_set();
            if (step != 1) {
                throw LSST_EXCEPT(pex::exceptions::InvalidParameterError,
                                  "Step for array Key indexing must be 1.");
            }
            return py::cast(self.slice(start, stop));
        } else {
            return py::cast(self[py::cast<int>(index)]);
        }
    });
    cls.def("slice", &KeyBase<Array<T>>::slice);
}
 
// Specializations for Key
 
template <typename T>
void declareKeyAccessors(PyKey<T> &cls) {
    cls.def("get", [](Key<T> const &self, BaseRecord &record) { return record.get(self); });
    cls.def("set", [](Key<T> const &self, BaseRecord &record, typename Key<T>::Value const &value) {
        record.set(self, value);
    });
}
 
template <typename U>
void declareKeyAccessors(PyKey<Array<U>> &cls) {
    auto getter = [](Key<Array<U>> const &self, BaseRecord &record) -> ndarray::Array<U, 1, 1> {
        return record[self];
    };
    auto setter = [](Key<Array<U>> const &self, BaseRecord &record, py::object const &value) {
        if (self.getSize() == 0) {
            // Variable-length array field: do a shallow copy, which requires a non-const
            // contiguous array.
            record.set(self, py::cast<ndarray::Array<U, 1, 1>>(value));
        } else {
            // Fixed-length array field: do a deep copy, which can work with a const
            // noncontiguous array.  But we need to check the size first, since the
            // penalty for getting that wrong is assert->abort.
            auto v = py::cast<ndarray::Array<U const, 1, 0>>(value);
            ndarray::ArrayRef<U, 1, 1> ref = record[self];
            if (v.size() != ref.size()) {
                throw LSST_EXCEPT(
                        pex::exceptions::LengthError,
                        (boost::format("Array sizes do not agree: %s != %s") % v.size() % ref.size()).str());
            }
            ref = v;
        }
        return;
    };
    cls.def("get", getter);
    cls.def("set", setter);
}
 
template <typename T>
void declareKeySpecializations(PyKey<T> &cls) {
    declareKeyAccessors(cls);
    cls.def_property_readonly("subfields", [](py::object const &) { return py::none(); });
    cls.def_property_readonly("subkeys", [](py::object const &) { return py::none(); });
    cls.def(py::pickle(
            [](Key<T> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getOffset());
            },
            [](py::tuple t) {
                int const NPARAMS = 1;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return detail::Access::makeKey<T>(t[0].cast<int>());
            }));
}
 
void declareKeySpecializations(PyKey<Flag> &cls) {
    declareKeyAccessors(cls);
    cls.def_property_readonly("subfields", [](py::object const &) { return py::none(); });
    cls.def_property_readonly("subkeys", [](py::object const &) { return py::none(); });
    cls.def("getBit", &Key<Flag>::getBit);
    cls.def(py::pickle(
            [](Key<Flag> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getOffset(), self.getBit());
            },
            [](py::tuple t) {
                int const NPARAMS = 2;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return detail::Access::makeKey(t[0].cast<int>(), t[1].cast<int>());
            }));
}
 
template <typename T>
void declareKeySpecializations(PyKey<Array<T>> &cls) {
    declareKeyAccessors(cls);
    cls.def_property_readonly("subfields", [](Key<Array<T>> const &self) -> py::object {
        py::list result;
        for (int i = 0; i < self.getSize(); ++i) {
            result.append(py::cast(i));
        }
        return py::tuple(result);
    });
    cls.def_property_readonly("subkeys", [](Key<Array<T>> const &self) -> py::object {
        py::list result;
        for (int i = 0; i < self.getSize(); ++i) {
            result.append(py::cast(self[i]));
        }
        return py::tuple(result);
    });
    cls.def(py::pickle(
            [](Key<Array<T>> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getOffset(), self.getElementCount());
            },
            [](py::tuple t) {
                int const NPARAMS = 2;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return detail::Access::makeKeyArray<T>(t[0].cast<int>(), t[1].cast<int>());
            }));
}
 
void declareKeySpecializations(PyKey<std::string> &cls) {
    declareKeyAccessors(cls);
    cls.def_property_readonly("subfields", [](py::object const &) { return py::none(); });
    cls.def_property_readonly("subkeys", [](py::object const &) { return py::none(); });
    cls.def(py::pickle(
            [](Key<std::string> const &self) {
                /* Return a tuple that fully encodes the state of the object */
                return py::make_tuple(self.getOffset(), self.getElementCount());
            },
            [](py::tuple t) {
                int const NPARAMS = 2;
                if (t.size() != NPARAMS) {
                    std::ostringstream os;
                    os << "Invalid number of parameters (" << t.size() << ") when unpickling; expected "
                       << NPARAMS;
                    throw std::runtime_error(os.str());
                }
                return detail::Access::makeKeyString(t[0].cast<int>(), t[1].cast<int>());
            }));
}
 
// Wrap all helper classes (FieldBase, KeyBase, Key, Field, SchemaItem) declarefor a Schema field type.
template <typename T>
void declareSchemaType(WrapperCollection &wrappers) {
    std::string suffix = FieldBase<T>::getTypeString();
    py::str pySuffix(suffix);
 
    py::object astropyUnit = py::module::import("astropy.units").attr("Unit");
 
    // FieldBase
    wrappers.wrapType(PyFieldBase<T>(wrappers.module, ("FieldBase" + suffix).c_str()),
                      [](auto &mod, auto &cls) {
                          cls.def_static("getTypeString", &FieldBase<T>::getTypeString);
                          declareFieldBaseSpecializations(cls);
                      });
 
    // KeyBase
    wrappers.wrapType(PyKeyBase<T>(wrappers.module, ("KeyBase" + suffix).c_str()), [](auto &mod, auto &cls) {
        cls.def_readonly_static("HAS_NAMED_SUBFIELDS", &KeyBase<T>::HAS_NAMED_SUBFIELDS);
        declareKeyBaseSpecializations(cls);
    });
 
    // Field
    wrappers.wrapType(PyField<T>(wrappers.module, ("Field" + suffix).c_str()), [pySuffix, astropyUnit](
                                                                                       auto &mod, auto &cls) {
        declareFieldSpecializations(cls);
 
        mod.attr("_Field")[pySuffix] = cls;
 
        cls.def(py::init([astropyUnit](  // capture by value to refcount in Python instead of dangle in C++
                                 std::string const &name, std::string const &doc, py::str const &units,
                                 py::object const &size, py::str const &parse_strict) {
                    astropyUnit(units, "parse_strict"_a = parse_strict);
                    std::string u = py::cast<std::string>(units);
                    if (size.is(py::none())) {
                        return new Field<T>(name, doc, u);
                    } else {
                        int s = py::cast<int>(size);
                        return new Field<T>(name, doc, u, s);
                    }
                }),
                "name"_a, "doc"_a = "", "units"_a = "", "size"_a = py::none(), "parse_strict"_a = "raise");
        cls.def("_addTo", [](Field<T> const &self, Schema &schema, bool doReplace) -> Key<T> {
            return schema.addField(self, doReplace);
        });
        cls.def("getName", &Field<T>::getName);
        cls.def("getDoc", &Field<T>::getDoc);
        cls.def("getUnits", &Field<T>::getUnits);
        cls.def("copyRenamed", &Field<T>::copyRenamed);
        utils::python::addOutputOp(cls, "__str__");
        utils::python::addOutputOp(cls, "__repr__");
    });
 
    // Key
    wrappers.wrapType(PyKey<T>(wrappers.module, ("Key" + suffix).c_str()), [pySuffix](auto &mod, auto &cls) {
        mod.attr("_Key")[pySuffix] = cls;
        cls.def(py::init<>());
        cls.def("__eq__", [](Key<T> const &self, Key<T> const &other) -> bool { return self == other; },
                py::is_operator());
        cls.def("__ne__", [](Key<T> const &self, Key<T> const &other) -> bool { return self != other; },
                py::is_operator());
        cls.def("isValid", &Key<T>::isValid);
        cls.def("getOffset", &Key<T>::getOffset);
        utils::python::addOutputOp(cls, "__str__");
        utils::python::addOutputOp(cls, "__repr__");
        // The Key methods below actually wrap templated methods on Schema and
        // SchemaMapper.  Rather than doing many-type overload resolution by
        // wrapping those methods directly, we use the visitor pattern by having
        // the wrappers for those methods delegate back to these non-templated
        // methods on the templated Key classes.
        cls.def("_findIn", [](Key<T> const &self, Schema const &schema) { return schema.find(self); });
        cls.def("_addMappingTo", [](Key<T> const &self, SchemaMapper &mapper, Field<T> const &field,
                                    bool doReplace) { return mapper.addMapping(self, field, doReplace); });
        cls.def("_addMappingTo", [](Key<T> const &self, SchemaMapper &mapper, std::string const &name,
                                    bool doReplace) { return mapper.addMapping(self, name, doReplace); });
        cls.def("_addMappingTo", [](Key<T> const &self, SchemaMapper &mapper, py::object const &,
                                    bool doReplace) { return mapper.addMapping(self, doReplace); });
        declareKeySpecializations(cls);
    });
 
    // SchemaItem
    wrappers.wrapType(PySchemaItem<T>(wrappers.module, ("SchemaItem" + suffix).c_str()),
                      [pySuffix](auto &mod, auto &cls) {
                          mod.attr("_SchemaItem")[pySuffix] = cls;
                          cls.def_readonly("key", &SchemaItem<T>::key);
                          cls.def_readonly("field", &SchemaItem<T>::field);
                          cls.def("getKey", [](SchemaItem<T> const &self) { return self.key; });
                          cls.def("getField", [](SchemaItem<T> const &self) { return self.field; });
                          cls.def("__getitem__", [](py::object const &self, int index) -> py::object {
                              if (index == 0) {
                                  return self.attr("key");
                              } else if (index == 1) {
                                  return self.attr("field");
                              }
                              // Have to raise IndexError not some LSST exception to get the
                              // right behavior when unpacking.
                              throw py::index_error("Index to SchemaItem must be 0 or 1.");
                          });
                          cls.def("__len__", [](py::object const &self) -> int { return 2; });
                          cls.def("__str__",
                                  [](py::object const &self) -> py::str { return py::str(py::tuple(self)); });
                          cls.def("__repr__", [](py::object const &self) -> py::str {
                              return py::str("SchemaItem(key={0.key}, field={0.field})").format(self);
                          });
                          cls.def(py::pickle(
                                  [](SchemaItem<T> const &self) {
                                      /* Return a tuple that fully encodes the state of the object */
                                      return py::make_tuple(self.key, self.field);
                                  },
                                  [](py::tuple t) {
                                      int const NPARAMS = 2;
                                      if (t.size() != NPARAMS) {
                                          std::ostringstream os;
                                          os << "Invalid number of parameters (" << t.size()
                                             << ") when unpickling; expected " << NPARAMS;
                                          throw std::runtime_error(os.str());
                                      }
                                      return SchemaItem<T>(t[0].cast<Key<T>>(), t[1].cast<Field<T>>());
                                  }));
                      });
}
 
// Helper class for Schema::find(name, func) that converts the result to Python.
// In C++14, this should be converted to a universal lambda.
struct MakePythonSchemaItem {
    template <typename T>
    void operator()(SchemaItem<T> const &item) {
        result = py::cast(item);
    }
 
    py::object result;
};
 
void declareSchema(WrapperCollection &wrappers) {
    wrappers.wrapType(PySchema(wrappers.module, "Schema"), [](auto &mod, auto &cls) {
        // wrap ComparisonFlags values as ints since we use them as bitflags,
        // not true enums
        cls.attr("EQUAL_KEYS") = py::cast(int(Schema::EQUAL_KEYS));
        cls.attr("EQUAL_NAMES") = py::cast(int(Schema::EQUAL_NAMES));
        cls.attr("EQUAL_DOCS") = py::cast(int(Schema::EQUAL_DOCS));
        cls.attr("EQUAL_UNITS") = py::cast(int(Schema::EQUAL_UNITS));
        cls.attr("EQUAL_FIELDS") = py::cast(int(Schema::EQUAL_FIELDS));
        cls.attr("EQUAL_ALIASES") = py::cast(int(Schema::EQUAL_ALIASES));
        cls.attr("IDENTICAL") = py::cast(int(Schema::IDENTICAL));
 
        cls.attr("VERSION") = py::cast(int(Schema::VERSION));
 
        cls.def(py::init<>());
        cls.def(py::init<Schema const &>());
        cls.def("__getitem__", [](Schema &self, std::string const &name) { return self[name]; });
        cls.def("__eq__", [](Schema const &self, Schema const &other) { return self == other; },
                py::is_operator());
        cls.def("__ne__", [](Schema const &self, Schema const &other) { return self != other; },
                py::is_operator());
        cls.def("getRecordSize", &Schema::getRecordSize);
        cls.def("getFieldCount", &Schema::getFieldCount);
        cls.def("getFlagFieldCount", &Schema::getFlagFieldCount);
        cls.def("getNonFlagFieldCount", &Schema::getNonFlagFieldCount);
        cls.def("find", [](py::object const &self, py::object const &key) -> py::object {
            try {
                if (py::isinstance<py::str>(key) || py::isinstance<py::bytes>(key)) {
                    Schema const &s = py::cast<Schema const &>(self);
                    std::string name = py::cast<std::string>(key);
                    MakePythonSchemaItem func;
                    s.findAndApply(name, func);
                    return func.result;
                }
                return key.attr("_findIn")(self);
            } catch (pex::exceptions::NotFoundError &err) {
                // Avoid API change by re-throwing as KeyError.
                PyErr_SetString(PyExc_KeyError, err.what());
                throw py::error_already_set();
            }
        });
        cls.def("getNames", &Schema::getNames, "topOnly"_a = false);
        cls.def("getAliasMap", &Schema::getAliasMap);
        cls.def("setAliasMap", &Schema::setAliasMap, "aliases"_a);
        cls.def("disconnectAliases", &Schema::disconnectAliases);
        cls.def("forEach", [](Schema &self, py::object &obj) { self.forEach(obj); });
        cls.def("compare", &Schema::compare, "other"_a, "flags"_a = int(Schema::EQUAL_KEYS));
        cls.def("contains", (int (Schema::*)(Schema const &, int) const) & Schema::contains, "other"_a,
                "flags"_a = int(Schema::EQUAL_KEYS));
        cls.def("__contains__", [](py::object const &self, py::object const &key) {
            try {
                self.attr("find")(key);
            } catch (py::error_already_set &err) {
                err.restore();
                PyErr_Clear();
                return false;
            }
            return true;
        });
        cls.def_static("readFits", (Schema(*)(std::string const &, int)) & Schema::readFits, "filename"_a,
                       "hdu"_a = fits::DEFAULT_HDU);
        cls.def_static("readFits", (Schema(*)(fits::MemFileManager &, int)) & Schema::readFits, "manager"_a,
                       "hdu"_a = fits::DEFAULT_HDU);
 
        cls.def("join",
                (std::string(Schema::*)(std::string const &, std::string const &) const) & Schema::join,
                "a"_a, "b"_a);
        cls.def("join",
                (std::string(Schema::*)(std::string const &, std::string const &, std::string const &)
                         const) &
                        Schema::join,
                "a"_a, "b"_a, "c"_a);
        cls.def("join",
                (std::string(Schema::*)(std::string const &, std::string const &, std::string const &,
                                        std::string const &) const) &
                        Schema::join,
                "a"_a, "b"_a, "c"_a, "d"_a);
        utils::python::addOutputOp(cls, "__str__");
        utils::python::addOutputOp(cls, "__repr__");
    });
}
 
void declareSubSchema(WrapperCollection &wrappers) {
    wrappers.wrapType(PySubSchema(wrappers.module, "SubSchema"), [](auto &mod, auto &cls) {
        cls.def("getNames", &SubSchema::getNames, "topOnly"_a = false);
        cls.def("getPrefix", &SubSchema::getPrefix);
        cls.def("asKey", [](SubSchema const &self) -> py::object {
            MakePythonSchemaItem func;
            self.apply(func);
            return func.result.attr("key");
        });
        cls.def("asField", [](SubSchema const &self) -> py::object {
            MakePythonSchemaItem func;
            self.apply(func);
            return func.result.attr("field");
        });
        cls.def("find", [](SubSchema const &self, std::string const &name) -> py::object {
            MakePythonSchemaItem func;
            self.findAndApply(name, func);
            return func.result;
        });
        cls.def("__getitem__", [](SubSchema &self, std::string const &name) { return self[name]; });
    });
}
 
}  // namespace
 
void wrapSchema(WrapperCollection &wrappers) {
    // We'll add instantiations of Field, Key, and SchemaItem to these private
    // dicts, and then in schemaContinued.py we'll add them to a TemplateMeta
    // ABC.
    auto &mod = wrappers.module;
    mod.attr("_Field") = py::dict();
    mod.attr("_Key") = py::dict();
    mod.attr("_SchemaItem") = py::dict();
 
    declareSchemaType<std::uint8_t>(wrappers);
    declareSchemaType<std::uint16_t>(wrappers);
    declareSchemaType<std::int32_t>(wrappers);
    declareSchemaType<std::int64_t>(wrappers);
    declareSchemaType<float>(wrappers);
    declareSchemaType<double>(wrappers);
    declareSchemaType<std::string>(wrappers);
    declareSchemaType<lsst::geom::Angle>(wrappers);
    declareSchemaType<Array<std::uint8_t>>(wrappers);
    declareSchemaType<Array<std::uint16_t>>(wrappers);
    declareSchemaType<Array<int>>(wrappers);
    declareSchemaType<Array<float>>(wrappers);
    declareSchemaType<Array<double>>(wrappers);
    declareSchemaType<Flag>(wrappers);
 
    declareSchema(wrappers);
    declareSubSchema(wrappers);
}
 
}  // namespace table
}  // namespace afw
}  // namespace lsst