From aa6205f6001f79f2f7045d08cf514ef2e0b4adff Mon Sep 17 00:00:00 2001
From: Pierre Marcenac <pierremarcenac@google.com>
Date: Tue, 4 Jun 2024 08:16:37 -0700
Subject: [PATCH] Support for python3.12

PiperOrigin-RevId: 640159257
---
 .github/workflows/pytest.yml                  |  2 +-
 setup.py                                      |  5 +++--
 .../core/features/audio_feature_test.py       |  4 ++++
 .../core/github_api/github_path.py            | 21 +++++++++++++++----
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index f66d4c9ea0a..df28115192c 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -84,7 +84,7 @@ jobs:
         num-shards: ${{ fromJson(needs.shards-job.outputs.num-shards) }}
         shard-id: ${{ fromJson(needs.shards-job.outputs.shard-ids) }}
         # TF suppported versions: https://www.tensorflow.org/install/pip#software_requirements
-        python-version: ['3.10', '3.11']
+        python-version: ['3.10', '3.11', '3.12']
         os-version: [ubuntu-latest]
 
     steps:
diff --git a/setup.py b/setup.py
index 76374426883..93670ccba42 100644
--- a/setup.py
+++ b/setup.py
@@ -65,7 +65,8 @@
     'array_record>=0.5.0;platform_system=="Linux"',
     'click',
     'dm-tree',
-    'etils[enp,epath,epy,etree]>=1.6.0',
+    'etils[enp,epath,epy,etree]>=1.6.0;python_version<"3.11"',
+    'etils[enp,epath,epy,etree]>=1.9.1;python_version>="3.11"',
     'immutabledict',
     'numpy',
     'promise',
@@ -99,7 +100,7 @@
     # 'tensorflow-docs @ git+https://github.com/tensorflow/docs#egg=tensorflow-docs',  # pylint: disable=line-too-long
     # Required by scripts/documentation/
     'pyyaml',
-    'tensorflow-io[tensorflow]',
+    'tensorflow-io[tensorflow];python_version<"3.12"',
 ]
 
 # Additional deps for formatting
diff --git a/tensorflow_datasets/core/features/audio_feature_test.py b/tensorflow_datasets/core/features/audio_feature_test.py
index 661b0d912da..89a35c8f9c4 100644
--- a/tensorflow_datasets/core/features/audio_feature_test.py
+++ b/tensorflow_datasets/core/features/audio_feature_test.py
@@ -110,6 +110,10 @@ def test_wav_file(self, dtype, lazy_decode, num_channels):
   @parameterized.product(lazy_decode=[True, False], num_channels=[1, 2, 8])
   def test_flac_file(self, lazy_decode, num_channels):
     if lazy_decode:
+      try:
+        import tensorflow_io  # pylint: disable=g-import-not-at-top, unused-import
+      except ImportError:
+        self.skipTest('`tensorflow_io` dependency is not available')
       if 'dev' in tf.__version__:
         self.skipTest('`tensorflow_io` is not compatible with `tf-nightly`')
 
diff --git a/tensorflow_datasets/core/github_api/github_path.py b/tensorflow_datasets/core/github_api/github_path.py
index 88093cfec1b..7ff0f8dc502 100644
--- a/tensorflow_datasets/core/github_api/github_path.py
+++ b/tensorflow_datasets/core/github_api/github_path.py
@@ -20,6 +20,7 @@
 import os
 import pathlib
 import posixpath
+import sys
 from typing import Iterator, Mapping, MutableMapping, Optional, Set, Tuple
 
 from etils import epath
@@ -217,10 +218,22 @@ class GithubPath(pathlib.PurePosixPath):
   ```
   """
 
-  def __new__(cls, *parts: epath.PathLike) -> 'GithubPath':
-    full_path = '/'.join(os.fspath(p) for p in parts)
-    _parse_github_path(full_path)
-    return super().__new__(cls, full_path.replace(_URI_PREFIX, '/github/', 1))
+  if sys.version_info < (3, 12):
+
+    def __new__(cls, *parts: epath.PathLike) -> 'GithubPath':
+      full_path = '/'.join(os.fspath(p) for p in parts)
+      _parse_github_path(full_path)
+      return super().__new__(cls, full_path.replace(_URI_PREFIX, '/github/', 1))
+
+  else:
+
+    def __init__(self, *parts: epath.PathLike):
+      full_path = '/'.join(os.fspath(p) for p in parts)
+      if full_path.startswith(_URI_PREFIX):
+        # If we already converted the prefix github:// -> /github, we don't do
+        # the check again.
+        _parse_github_path(full_path)
+      super().__init__(full_path.replace(_URI_PREFIX, '/github/', 1))
 
   @functools.cached_property
   def _path_str(self) -> str: