add py binding

This commit is contained in:
Shunsuke Kanda 2021-07-10 22:42:36 +09:00
parent 1e9ad82236
commit 632b74a380
2 changed files with 251 additions and 0 deletions

114
pybind/setup.py Executable file
View file

@ -0,0 +1,114 @@
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
import sys
import setuptools
__version__ = '0.1.0'
class get_pybind_include(object):
"""Helper class to determine the pybind11 include path
The purpose of this class is to postpone importing pybind11
until it is actually installed, so that the ``get_include()``
method can be invoked. """
def __init__(self, user=False):
self.user = user
def __str__(self):
import pybind11
return pybind11.get_include(self.user)
ext_modules = [
Extension(
'xcdat',
['xcdat_pybind.cpp'],
include_dirs=[
# Path to pybind11 headers
get_pybind_include(),
get_pybind_include(user=True),
],
language='c++'
),
]
# As of Python 3.6, CCompiler has a `has_flag` method.
# cf http://bugs.python.org/issue26689
def has_flag(compiler, flagname):
"""Return a boolean indicating whether a flag name is supported on
the specified compiler.
"""
import tempfile
with tempfile.NamedTemporaryFile('w', suffix='.cpp') as f:
f.write('int main (int argc, char **argv) { return 0; }')
try:
compiler.compile([f.name], extra_postargs=[flagname])
except setuptools.distutils.errors.CompileError:
return False
return True
def cpp_flag(compiler):
"""Return the -std=c++[17] compiler flag.
The newer version is prefered over c++17 (when it is available).
"""
flags = ['-std=c++17']
for flag in flags:
if has_flag(compiler, flag):
return flag
raise RuntimeError(
'Unsupported compiler -- at least C++17 support is needed!')
class BuildExt(build_ext):
"""A custom build extension for adding compiler-specific options."""
c_opts = {
'msvc': ['/EHsc'],
'unix': ['-DNDEBUG', '-march=native', '-O3'],
}
l_opts = {
'msvc': [],
'unix': [],
}
if sys.platform == 'darwin':
darwin_opts = ['-stdlib=libc++', '-mmacosx-version-min=10.7']
c_opts['unix'] += darwin_opts
l_opts['unix'] += darwin_opts
def build_extensions(self):
ct = self.compiler.compiler_type
opts = self.c_opts.get(ct, [])
link_opts = self.l_opts.get(ct, [])
if ct == 'unix':
opts.append('-DVERSION_INFO="%s"' %
self.distribution.get_version())
opts.append(cpp_flag(self.compiler))
if has_flag(self.compiler, '-fvisibility=hidden'):
opts.append('-fvisibility=hidden')
elif ct == 'msvc':
opts.append('/DVERSION_INFO=\\"%s\\"' %
self.distribution.get_version())
for ext in self.extensions:
ext.extra_compile_args = opts
ext.extra_link_args = link_opts
build_ext.build_extensions(self)
setup(
name='xcdat',
version=__version__,
author='Kampersanda',
description='Xcdat: Fast compressed trie dictionary',
long_description='',
ext_modules=ext_modules,
install_requires=['pybind11>=2.4'],
setup_requires=['pybind11>=2.4'],
cmdclass={'build_ext': BuildExt},
zip_safe=False,
)

137
pybind/xcdat_pybind.cpp Normal file
View file

@ -0,0 +1,137 @@
#include <memory>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <xcdat.hpp>
namespace {
class Keyset {
private:
std::vector<std::string> m_keys;
public:
Keyset() = default;
virtual ~Keyset() = default;
void append(const std::string& key) {
m_keys.push_back(key);
}
void complete() {
std::sort(m_keys.begin(), m_keys.end());
m_keys.erase(std::unique(m_keys.begin(), m_keys.end()), m_keys.end());
}
const std::vector<std::string>& get() const {
return m_keys;
}
};
} // namespace
namespace py = pybind11;
template <class Trie>
void define_Trie(py::module_& m, const std::string& suffix) {
const std::string name = "Trie" + suffix;
py::class_<Trie>(m, name.c_str(), "A class of trie dictionary") //
.def(py::init([](const Keyset& keyset) { //
return std::make_unique<Trie>(keyset.get());
}),
"Build the trie dictionary from the keyset.", py::arg("keyset"))
.def(py::init([](const std::string& filepath) { //
return std::make_unique<Trie>(xcdat::load<Trie>(filepath));
}),
"Load the trie dictionary from the filepath.", py::arg("filepath"))
.def(
"save", [](const Trie& self, const std::string& filepath) { xcdat::save(self, filepath); },
"Save the trie dictionary to the filepath", py::arg("filepath"))
.def(
"memory_in_bytes", [](const Trie& self) { return xcdat::memory_in_bytes(self); },
"Get the memory usage of the trie dictionary")
.def("bin_mode", &Trie::bin_mode, "Check if the binary mode.")
.def("num_keys", &Trie::num_keys, "Get the number of stored keywords.")
.def("alphabet_size", &Trie::alphabet_size, "Get the alphabet size.")
.def("max_length", &Trie::max_length, "Get the maximum length of keywords.")
.def("num_nodes", &Trie::num_nodes, "Get the number of trie nodes.")
.def("num_units", &Trie::num_units, "Get the number of DA units.")
.def("num_free_units", &Trie::num_free_units, "Get the number of unused DA units.")
.def("tail_length", &Trie::tail_length, "Get the length of TAIL vector.")
.def("lookup", &Trie::lookup, "Lookup the ID of the keyword.", py::arg("key"))
.def(
"decode", [](const Trie& self, std::uint64_t id) { return self.decode(id); },
"Decode the keyword associated with the ID.", py::arg("id"))
.def(
"make_prefix_iterator",
[](const Trie& self, const std::string& key) {
return std::make_unique<typename Trie::prefix_iterator>(self.make_prefix_iterator(key));
},
"Make the common prefix searcher for the given keyword.", py::arg("key"))
.def(
"make_predictive_iterator",
[](const Trie& self, const std::string& key) {
return std::make_unique<typename Trie::predictive_iterator>(self.make_predictive_iterator(key));
},
"Make the predictive searcher for the keyword.", py::arg("key"))
.def(
"make_enumerative_iterator",
[](const Trie& self) {
return std::make_unique<typename Trie::enumerative_iterator>(self.make_enumerative_iterator());
},
"Make the enumerator.");
}
template <class Trie>
void define_PrefixIterator(py::module_& m, const std::string& suffix) {
const std::string name = "PrefixIterator" + suffix;
py::class_<typename Trie::prefix_iterator>(
m, name.c_str(),
"An iterator class for common prefix search.\n"
"It enumerates all the keywords contained as prefixes of a given string.\n"
"It should be instantiated via the function 'make_prefix_iterator'.")
.def("next", &Trie::prefix_iterator::next,
"Increment the iterator. Return false if the iteration is terminated.")
.def("id", &Trie::prefix_iterator::id, "Get the result ID.")
.def("decoded", &Trie::prefix_iterator::decoded, "Get the result keyword.");
}
template <class Trie>
void define_PredictiveIterator(py::module_& m, const std::string& suffix) {
const std::string name = "PredictiveIterator" + suffix;
py::class_<typename Trie::predictive_iterator>(
m, name.c_str(),
"An iterator class for predictive search or enumeration.\n"
"It enumerates all the keywords starting with a given string (including an empty string).\n"
"It should be instantiated via the function 'make_prefix_iterator' or 'make_enumerative_iterator'.") //
.def("next", &Trie::predictive_iterator::next,
"Increment the iterator. Return false if the iteration is terminated.")
.def("id", &Trie::predictive_iterator::id, "Get the result ID.")
.def("decoded", &Trie::predictive_iterator::decoded, "Get the result keyword.");
}
PYBIND11_MODULE(xcdat, m) {
m.doc() = "Xcdat: Fast compressed trie dictionary";
py::class_<Keyset>(m, "Keyset", "A class to store keywords for building xcdat.Trie.") //
.def(py::init<>())
.def("append", &Keyset::append, "Append the given key.", py::arg("key"))
.def("complete", &Keyset::complete, "Complete the keyset to build a xcdat.Trie.");
define_Trie<xcdat::trie_8_type>(m, "8");
define_PrefixIterator<xcdat::trie_8_type>(m, "8");
define_PredictiveIterator<xcdat::trie_8_type>(m, "8");
define_Trie<xcdat::trie_16_type>(m, "16");
define_PrefixIterator<xcdat::trie_16_type>(m, "16");
define_PredictiveIterator<xcdat::trie_16_type>(m, "16");
define_Trie<xcdat::trie_7_type>(m, "7");
define_PrefixIterator<xcdat::trie_7_type>(m, "7");
define_PredictiveIterator<xcdat::trie_7_type>(m, "7");
define_Trie<xcdat::trie_15_type>(m, "15");
define_PrefixIterator<xcdat::trie_15_type>(m, "15");
define_PredictiveIterator<xcdat::trie_15_type>(m, "15");
}