Migrate package metadata to pyproject.toml.

.bazelversion

@@ -0,0 +1 @@
+7.7.0

kokoro/github/ubuntu/cpu/build.sh

@@ -16,20 +16,35 @@
 set -e
 set -x
 
-cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
+PYENV_ROOT="/home/kbuilder/.pyenv"
+PYTHON_VERSION=${PYTHON_VERSION:-"3.9"}
 
-sudo update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1
+echo "Installing pyenv.."
+git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+export PATH="/home/kbuilder/.local/bin:$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
 
-PYTHON_BINARY="/usr/bin/python3.9"
-PIP_TEST_PREFIX=bazel_pip
+echo "Python setup..."
+pyenv install -s "$PYTHON_VERSION"
+pyenv global "$PYTHON_VERSION"
 
-"${PYTHON_BINARY}" -m venv venv
-source venv/bin/activate
+cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
+
+PIP_TEST_PREFIX=bazel_pip
 
 # Check the python version
 python --version
 python3 --version
 
+python -m venv build_venv
+source build_venv/bin/activate
+
+# update pip
+pip install --upgrade pip
+
+# Install build
+pip install build
+
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
@@ -43,14 +58,17 @@ if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}*
 fi
 
 bazel clean
-pip install -r requirements-dev.txt --progress-bar off
-pip install tf-keras-nightly tf-nightly --progress-bar off --upgrade
+pip install --group test-nightly --progress-bar off --upgrade
 # We need to remove the dependency on tensorflow to test nightly
-# The dependencies will be provided by tf-nightly
-perl  -i -lpe '$k+= s/tensorflow>=2\.[0-9]+\.[0-9]+(,<=?[0-9.]+)?;/tf-nightly;/g; END{exit($k != 1)}' setup.py
-python3 setup.py bdist_wheel
-pip uninstall -y tensorflow_gnn
+python3 -m build --wheel
+deactivate
+
+# Start the test environment.
+python3 -m venv test_venv
+source test_venv/bin/activate
+pip install --upgrade pip
 pip install dist/tensorflow_gnn-*.whl
+pip install --group test --progress-bar off --upgrade
 
 echo "Final packages after all pip commands:"
 pip list

kokoro/github/ubuntu/cpu/build_versioned.sh

@@ -16,24 +16,35 @@
 set -e
 set -x
 
-cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
+PYENV_ROOT="/home/kbuilder/.pyenv"
+PYTHON_VERSION=${PYTHON_VERSION:-"3.9"}
+
+echo "Installing pyenv.."
+git clone https://github.com/pyenv/pyenv.git "$PYENV_ROOT"
+export PATH="/home/kbuilder/.local/bin:$PYENV_ROOT/bin:$PATH"
+eval "$(pyenv init --path)"
 
-sudo apt-get install -y "python${PYTHON_VERSION}"
+echo "Python setup..."
+pyenv install -s "$PYTHON_VERSION"
+pyenv global "$PYTHON_VERSION"
 
-# Update alternatives, taken from corresponding Keras OSS test script
-sudo update-alternatives --install /usr/bin/python3 python3 "/usr/bin/python$PYTHON_VERSION" 1
-sudo apt-get install -y python$PYTHON_VERSION-venv
+cd "${KOKORO_ARTIFACTS_DIR}/github/gnn/"
 
-PYTHON_BINARY="/usr/bin/python${PYTHON_VERSION}"
 PIP_TEST_PREFIX=bazel_pip
 
-"${PYTHON_BINARY}" -m venv venv
-source venv/bin/activate
+python -m venv build_venv
+source build_venv/bin/activate
 
 # Debug messages to indicate the python version
 python --version
 python3 --version
 
+# update pip
+pip install --upgrade pip
+
+# Install build
+pip install build
+
 TEST_ROOT=$(pwd)/${PIP_TEST_PREFIX}
 rm -rf "$TEST_ROOT"
 mkdir -p "$TEST_ROOT"
@@ -54,16 +65,22 @@ if [[ -n "${USE_BAZEL_VERSION}" && $(bazel --version) != *${USE_BAZEL_VERSION}*
 fi
 
 bazel clean
-pip install -r requirements-dev.txt --progress-bar off
 pip install tensorflow=="${TF_VERSION}" --progress-bar off --upgrade
 if [[ "$TF_USE_LEGACY_KERAS" == 1 ]]; then
   pip install tf-keras=="${TF_VERSION}" --progress-bar off --upgrade
 fi
-python3 setup.py bdist_wheel
-pip uninstall -y tensorflow_gnn
+python3 -m build --wheel
+deactivate
+
+# Start the test environment.
+python3 -m venv test_venv
+source test_venv/bin/activate
+pip install --upgrade pip
+pip install --group test --progress-bar off --upgrade
+
 pip install dist/tensorflow_gnn-*.whl
 
 echo "Final packages after all pip commands:"
 pip list
 
-bazel test --test_env=TF_USE_LEGACY_KERAS --build_tag_filters="${tag_filters}" --test_tag_filters="${tag_filters}" --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...
+bazel test --test_env="TF_USE_LEGACY_KERAS=1" --build_tag_filters="${tag_filters}" --test_tag_filters="${tag_filters}" --test_output=errors --verbose_failures=true --build_tests_only --define=no_tfgnn_py_deps=true --keep_going --experimental_repo_remote_exec //bazel_pip/tensorflow_gnn/...

pyproject.toml

@@ -0,0 +1,79 @@
+[project]
+name = "tensorflow_gnn"
+description = "A library for building scalable graph neural networks in TensorFlow."
+readme = "README.md"
+keywords = ["tensorflow", "gnn", "graph"]
+dynamic = ["version"]
+license = "Apache-2.0"
+requires-python = ">=3.9"
+dependencies = [
+  "google-vizier>=0.0.13, !=0.1.23, !=0.1.24",  # b/394062744
+  "ml-collections",
+  "networkx",
+  "pyarrow",
+  "tensorflow>=2.12.0, <3; platform_machine != 'arm64' or platform_system != 'Darwin'",
+  "tensorflow-macos>=2.12.0, <3; platform_machine == 'arm64' and platform_system == 'Darwin'",
+  "apache-beam>=2.54",
+]
+authors = [
+  {name = "Google LLC", email = "tensorflow-gnn@googlegroups.com"},
+]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "Operating System :: MacOS :: MacOS X",
+  "Operating System :: POSIX :: Linux",
+  "Operating System :: Microsoft :: Windows",
+  "Programming Language :: Python",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3 :: Only",
+  "Topic :: Scientific/Engineering",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Scientific/Engineering :: Mathematics",
+  "Topic :: Software Development",
+  "Topic :: Software Development :: Libraries",
+  "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+[project.urls]
+homepage = "https://github.com/tensorflow/gnn"
+source = "https://github.com/tensorflow/gnn.git"
+
+[dependency-groups]
+test = [
+  "mock",
+  "tf-keras",
+]
+test-nightly = [
+  "tf-nightly",
+  "tf-keras-nightly",
+  {include-group = "test"},
+]
+
+[project.scripts]
+tfgnn_convert_ogb_dataset = 'tensorflow_gnn.converters.ogb.convert_ogb_dataset:main'
+tfgnn_sampler = 'tensorflow_gnn.experimental.sampler.beam.sampler:main'
+# copybara:comment_begin(NetworkX utils)
+tfgnn_nx_converter = 'tensorflow_gnn.sampler.nx_converter:main'
+tfgnn_nx_generator = 'tensorflow_gnn.sampler.nx_generator:main'
+# copybara:comment_end
+tfgnn_generate_training_data = 'tensorflow_gnn.tools.generate_training_data:main'
+tfgnn_print_training_data = 'tensorflow_gnn.tools.print_training_data:main'
+tfgnn_sampled_stats = 'tensorflow_gnn.tools.sampled_stats:main'
+tfgnn_validate_graph_schema = 'tensorflow_gnn.tools.validate_graph_schema:main'
+
+[build-system]
+requires = ["setuptools", "numpy"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.dynamic]
+version = {attr = "tensorflow_gnn.version.__version__"}
+
+[tool.setuptools.packages.find]
+exclude = ["examples*", "package*", "testdata*", "build*", "bazel*"]
+
+[tool.setuptools.package-data]
+tensorflow_gnn = ["*.proto"]
+

setup.py

@@ -16,20 +16,13 @@
 
 import os
 import platform
+import shutil
 import subprocess
 import sys
-
 import setuptools
-from setuptools import find_namespace_packages
-from setuptools import setup
-from setuptools.command.install import install
-from setuptools.dist import Distribution
-# pylint:disable=g-bad-import-order
-# setuptools must be imported prior to distutils.
-# TODO(b/317036825): Stop using distutils, Python 3.12 doesn't have it.
-from distutils import spawn
-from distutils.command import build
-# pylint:enable=g-bad-import-order
+from setuptools import dist
+from setuptools.command import build
+from setuptools.command import install
 
 
 class _BuildCommand(build.build):
@@ -57,7 +50,8 @@ def _build_cc_extensions(self):
   # sequentially when running a "build" command, if the second item in the tuple
   # (predicate method) is evaluated to true.
   sub_commands = [
-      ('bazel_build', _build_cc_extensions)] + build.build.sub_commands
+      ('bazel_build', _build_cc_extensions)
+  ] + build.build.sub_commands
 
 
 class _BazelBuildCommand(setuptools.Command):
@@ -71,7 +65,7 @@ def initialize_options(self):
     pass
 
   def finalize_options(self):
-    self._bazel_cmd = spawn.find_executable('bazel')
+    self._bazel_cmd = shutil.which('bazel')
     if not self._bazel_cmd:
       raise RuntimeError(
           'Could not find "bazel" binary. Please visit '
@@ -98,14 +92,14 @@ def run(self):
 # TFDV is not a purelib. However because of the extension module is not built
 # by setuptools, it will be incorrectly treated as a purelib. The following
 # works around that bug.
-class _InstallPlatlibCommand(install):
+class _InstallPlatlibCommand(install.install):
 
   def finalize_options(self):
-    install.finalize_options(self)
+    install.install.finalize_options(self)
     self.install_lib = self.install_platlib
 
 
-class _SourceDistributionWithProtos(Distribution):
+class _SourceDistributionWithProtos(dist.Distribution):
   """Proto-only deps don't require OS specific wheels."""
 
   def is_pure(self):
@@ -115,95 +109,11 @@ def has_ext_modules(self):
     return False
 
 
-def get_version():
-  """Get version from version module."""
-  version_path = os.path.join(os.path.dirname(__file__), 'tensorflow_gnn')
-  sys.path.insert(0, version_path)
-  # pytype: disable=import-error  # pylint: disable=g-import-not-at-top
-  from version import __version__ as v
-  return v
-
-
-# Get the long description from the README file.
-# TODO(b/316329189): The relative links from README.md don't work on PyPI.
-with open('README.md') as fp:
-  _LONG_DESCRIPTION = fp.read()
-
-
-console_scripts = [
-    'tensorflow_gnn.converters.ogb.convert_ogb_dataset',
-    'tensorflow_gnn.experimental.sampler.beam.sampler',
-    # copybara:uncomment_begin(NetworkX utils)
-    # 'tensorflow_gnn.sampler.nx_converter',
-    # 'tensorflow_gnn.sampler.nx_generator',
-    # copybara:uncomment_end
-    'tensorflow_gnn.tools.generate_training_data',
-    'tensorflow_gnn.tools.print_training_data',
-    'tensorflow_gnn.tools.sampled_stats',
-    'tensorflow_gnn.tools.validate_graph_schema',
-]
-
-
-setup(
-    name='tensorflow-gnn',
-    version=get_version(),
-    author='Google LLC',
-    author_email='tensorflow-gnn@googlegroups.com',
-    license='Apache 2.0',
-    classifiers=[
-        'Development Status :: 5 - Production/Stable',
-        'Intended Audience :: Developers',
-        'Intended Audience :: Science/Research',
-        'License :: OSI Approved :: Apache Software License',
-        'Operating System :: MacOS :: MacOS X',
-        'Operating System :: POSIX :: Linux',
-        'Operating System :: Microsoft :: Windows',
-        'Programming Language :: Python',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3 :: Only',
-        'Topic :: Scientific/Engineering',
-        'Topic :: Scientific/Engineering :: Artificial Intelligence',
-        'Topic :: Scientific/Engineering :: Mathematics',
-        'Topic :: Software Development',
-        'Topic :: Software Development :: Libraries',
-        'Topic :: Software Development :: Libraries :: Python Modules',
-    ],
-    namespace_packages=[],
-    install_requires=[
-        'google-vizier>=0.0.13,!=0.1.23,!=0.1.24',  # b/394062744
-        'ml-collections',
-        'networkx',
-        'pyarrow',
-        # pylint:disable=g-line-too-long
-        'tensorflow>=2.12.0,<3; platform_machine != "arm64" or platform_system != "Darwin"',
-        'tensorflow-macos>=2.12.0,<3; platform_machine == "arm64" and platform_system == "Darwin"',
-        # pylint:enable=g-line-too-long
-        'apache-beam',
-    ],
-    python_requires='>=3.9,<4',
-    packages=find_namespace_packages(
-        exclude=['examples*', 'package*', 'testdata*'],
-    ),
-    include_package_data=True,
-    package_data={'': ['*.proto']},
-    zip_safe=False,
+setuptools.setup(
     distclass=_SourceDistributionWithProtos,
-    description='A library for building scalable graph neural networks in TensorFlow.',
-    long_description=_LONG_DESCRIPTION,
-    long_description_content_type='text/markdown',
-    keywords='tensorflow gnn graph',
-    url='https://github.com/tensorflow/gnn',
-    download_url='https://github.com/tensorflow/gnn.git',
-    requires=[],
     cmdclass={
         'install': _InstallPlatlibCommand,
         'build': _BuildCommand,
         'bazel_build': _BazelBuildCommand,
     },
-    entry_points={
-        'console_scripts': [
-            'tfgnn_{}={}:main'.format(libname.split('.')[-1], libname)
-            for libname in console_scripts
-        ],
-    }
 )

tensorflow_gnn/graph/tag_utils_test.py

@@ -30,6 +30,7 @@ class ReverseTagTest(parameterized.TestCase):
       ("Source", const.SOURCE, const.TARGET),
       ("Target", const.TARGET, const.SOURCE))
   def test(self, tag, expected):
+    assert True
     actual = tag_utils.reverse_tag(tag)
     self.assertEqual(expected, actual)