From: Michael Fladischer <FladischerMichael@fladi.at>
Date: Mon, 31 Jul 2023 21:13:42 +0000
Subject: Simplify setup.py by using pybind11.setup_helpers.Pybind11Extension.

---
 setup.py | 112 ++++++++++++---------------------------------------------------
 1 file changed, 21 insertions(+), 91 deletions(-)

diff --git a/setup.py b/setup.py
index e04d77d..2e146cc 100644
--- a/setup.py
+++ b/setup.py
@@ -15,21 +15,14 @@
 """Setup for pip package."""
 
 import os
-import platform
-import shutil
-import subprocess
-import sys
-import sysconfig
 
-import setuptools
-from setuptools.command import build_ext
-
-here = os.path.dirname(os.path.abspath(__file__))
+from setuptools import find_packages, setup
+from pybind11.setup_helpers import Pybind11Extension, build_ext
 
 
 def _get_tree_version():
   """Parse the version string from tree/__init__.py."""
-  with open(os.path.join(here, 'tree', '__init__.py')) as f:
+  with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tree', '__init__.py')) as f:
     try:
       version_line = next(line for line in f if line.startswith('__version__'))
     except StopIteration:
@@ -40,93 +33,17 @@ def _get_tree_version():
       return ns['__version__']
 
 
-class CMakeExtension(setuptools.Extension):
-  """An extension with no sources.
-
-  We do not want distutils to handle any of the compilation (instead we rely
-  on CMake), so we always pass an empty list to the constructor.
-  """
-
-  def __init__(self, name, source_dir=''):
-    super().__init__(name, sources=[])
-    self.source_dir = os.path.abspath(source_dir)
-
-
-class BuildCMakeExtension(build_ext.build_ext):
-  """Our custom build_ext command.
-
-  Uses CMake to build extensions instead of a bare compiler (e.g. gcc, clang).
-  """
-
-  def run(self):
-    self._check_build_environment()
-    for ext in self.extensions:
-      self.build_extension(ext)
-
-  def _check_build_environment(self):
-    """Check for required build tools: CMake, C++ compiler, and python dev."""
-    try:
-      subprocess.check_call(['cmake', '--version'])
-    except OSError as e:
-      ext_names = ', '.join(e.name for e in self.extensions)
-      raise RuntimeError(
-          f'CMake must be installed to build the following extensions: {ext_names}'
-      ) from e
-    print('Found CMake')
-
-  def build_extension(self, ext):
-    extension_dir = os.path.abspath(
-        os.path.dirname(self.get_ext_fullpath(ext.name)))
-    build_cfg = 'Debug' if self.debug else 'Release'
-    cmake_args = [
-        f'-DPython3_ROOT_DIR={sys.prefix}',
-        f'-DPython3_EXECUTABLE={sys.executable}',
-        f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extension_dir}',
-        f'-DCMAKE_BUILD_TYPE={build_cfg}'
-    ]
-    if platform.system() != 'Windows':
-      cmake_args.extend([
-          f'-DPython3_LIBRARY={sysconfig.get_paths()["stdlib"]}',
-          f'-DPython3_INCLUDE_DIR={sysconfig.get_paths()["include"]}',
-      ])
-    if platform.system() == 'Darwin' and os.environ.get('ARCHFLAGS'):
-      osx_archs = []
-      if '-arch x86_64' in os.environ['ARCHFLAGS']:
-        osx_archs.append('x86_64')
-      if '-arch arm64' in os.environ['ARCHFLAGS']:
-        osx_archs.append('arm64')
-      cmake_args.append(f'-DCMAKE_OSX_ARCHITECTURES={";".join(osx_archs)}')
-    os.makedirs(self.build_temp, exist_ok=True)
-    subprocess.check_call(
-        ['cmake', '-S', ext.source_dir, '-B', self.build_temp] + cmake_args)
-    num_jobs = ()
-    if self.parallel:
-      num_jobs = (f'-j{self.parallel}',)
-    subprocess.check_call([
-        'cmake', '--build', self.build_temp, *num_jobs, '--config', build_cfg
-    ])
-
-    # Force output to <extension_dir>/. Amends CMake multigenerator output paths
-    # on Windows and avoids Debug/ and Release/ subdirs, which is CMake default.
-    tree_dir = os.path.join(extension_dir, 'tree')  # pylint:disable=unreachable
-    for cfg in ('Release', 'Debug'):
-      cfg_dir = os.path.join(extension_dir, cfg)
-      if os.path.isdir(cfg_dir):
-        for f in os.listdir(cfg_dir):
-          shutil.move(os.path.join(cfg_dir, f), tree_dir)
-
-
-setuptools.setup(
+setup(
     name='dm-tree',
     version=_get_tree_version(),
     url='https://github.com/deepmind/tree',
     description='Tree is a library for working with nested data structures.',
     author='DeepMind',
     author_email='tree-copybara@google.com',
-    long_description=open(os.path.join(here, 'README.md')).read(),
+    long_description=open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'README.md')).read(),
     long_description_content_type='text/markdown',
     # Contained modules and scripts.
-    packages=setuptools.find_packages(),
+    packages=find_packages(),
     python_requires='>=3.10',
     install_requires=[
         'absl-py>=0.6.1',
@@ -140,8 +57,21 @@ setuptools.setup(
         'wrapt>=1.11.2',
     ],
     test_suite='tree',
-    cmdclass=dict(build_ext=BuildCMakeExtension),
-    ext_modules=[CMakeExtension('_tree', source_dir='tree')],
+    cmdclass={"build_ext": build_ext},
+    ext_modules=[
+        Pybind11Extension(
+            'tree._tree',
+            sources=['tree/tree.cc'],
+            libraries=[
+                'absl_int128',
+                'absl_raw_hash_set',
+                'absl_raw_logging_internal',
+                'absl_strings',
+                'absl_synchronization',
+                'absl_throw_delegate',
+            ]
+        )
+    ],
     zip_safe=False,
     # PyPI package information.
     classifiers=[