bernsteinbear.com
this repo has no description
1#!/usr/bin/env python3
2# Copyright (c) Facebook, Inc. and its affiliates. (http://www.facebook.com)
3import _imp
4import distutils.sysconfig
5import importlib.machinery
6import os
7import tempfile
8import unittest
9from distutils.core import setup
10from distutils.dist import Distribution
11from distutils.extension import Extension
12
13
14class DistutilsModuleTest(unittest.TestCase):
15 def test_sysconfig_get_python_inc_returns_string(self):
16 python_inc = distutils.sysconfig.get_python_inc()
17 self.assertTrue(os.path.isdir(python_inc))
18 self.assertTrue(os.path.isfile(os.path.join(python_inc, "Python.h")))
19
20 def test_sysconfig_get_python_lib_returns_string(self):
21 python_lib = distutils.sysconfig.get_python_lib()
22 self.assertIn("site-packages", python_lib)
23
24 def test_build_ext_with_c_file_creates_library(self):
25 with tempfile.TemporaryDirectory() as dir_path:
26 source_dir = f"{dir_path}/src"
27 build_dir = f"{dir_path}/build"
28 os.makedirs(source_dir)
29 os.makedirs(build_dir)
30 file_path = f"{source_dir}/foo.c"
31 with open(file_path, "w") as fp:
32 fp.write(
33 """\
34#include "Python.h"
35PyObject *func(PyObject *a0, PyObject *a1) {
36 (void)a0;
37 (void)a1;
38 return PyUnicode_FromString("and now for something completely different");
39}
40static PyMethodDef methods[] = {
41 {"announce", func, METH_NOARGS, "bar docu"},
42 {NULL, NULL},
43};
44static PyModuleDef def = {
45 PyModuleDef_HEAD_INIT,
46 "foo",
47 NULL,
48 0,
49 methods,
50};
51PyMODINIT_FUNC PyInit_foo(void) {
52 return PyModule_Create(&def);
53}
54"""
55 )
56 self.assertEqual(len(os.listdir(source_dir)), 1)
57
58 # Setup compilation settings
59 ext = Extension(name="foo", sources=[file_path])
60 self.assertIsInstance(ext, Extension)
61
62 # Compile C File
63 dist = setup(
64 name="test_so_compilation",
65 ext_modules=[ext],
66 script_args=["--quiet", "build_ext", "--build-temp", build_dir],
67 options={"build_ext": {"build_lib": source_dir}},
68 )
69 self.assertIsInstance(dist, Distribution)
70
71 # Check directory contents
72 dir_contents = sorted(os.listdir(source_dir))
73 self.assertEqual(len(dir_contents), 2)
74 self.assertTrue(dir_contents[0].endswith(".c"))
75 self.assertTrue(dir_contents[1].endswith(".so"))
76
77 lib_file = f"{source_dir}/{dir_contents[1]}"
78 spec = importlib.machinery.ModuleSpec("foo", None, origin=lib_file)
79 module = _imp.create_dynamic(spec)
80 self.assertEqual(
81 module.announce(), "and now for something completely different"
82 )
83
84 def test_build_ext_with_cpp_file_creates_library(self):
85 with tempfile.TemporaryDirectory() as dir_path:
86 source_dir = f"{dir_path}/src"
87 build_dir = f"{dir_path}/build"
88 os.makedirs(source_dir)
89 os.makedirs(build_dir)
90 file_path = f"{source_dir}/foo.cpp"
91 with open(file_path, "w") as fp:
92 fp.write(
93 """\
94#include <sstream>
95#include "Python.h"
96PyObject* func(PyObject*, PyObject*) {
97 std::stringstream ss;
98 ss << "Hello" << ' ' << "world";
99 return PyUnicode_FromString(ss.str().c_str());
100}
101static PyMethodDef methods[] = {
102 {"greet", func, METH_NOARGS, "bar docu"},
103 {nullptr, nullptr},
104};
105static PyModuleDef def = {
106 PyModuleDef_HEAD_INIT,
107 "foo",
108 nullptr,
109 0,
110 methods,
111};
112PyMODINIT_FUNC PyInit_foo() {
113 return PyModule_Create(&def);
114}
115"""
116 )
117 self.assertEqual(len(os.listdir(source_dir)), 1)
118
119 # Setup compilation settings
120 ext = Extension(name="foo", sources=[file_path])
121 self.assertIsInstance(ext, Extension)
122
123 dist = setup(
124 name="test_so_compilation",
125 ext_modules=[ext],
126 script_args=["--quiet", "build_ext", "--build-temp", build_dir],
127 options={"build_ext": {"build_lib": source_dir}},
128 )
129 self.assertIsInstance(dist, Distribution)
130
131 # Check directory contents
132 dir_contents = sorted(os.listdir(source_dir))
133 self.assertEqual(len(dir_contents), 2)
134 self.assertTrue(dir_contents[0].endswith(".cpp"))
135 self.assertTrue(dir_contents[1].endswith(".so"))
136
137 lib_file = f"{source_dir}/{dir_contents[1]}"
138 spec = importlib.machinery.ModuleSpec("foo", None, origin=lib_file)
139 module = _imp.create_dynamic(spec)
140 self.assertEqual(module.greet(), "Hello world")
141
142
143if __name__ == "__main__":
144 unittest.main()