From 9a2f2b68dba1f3aa417ef22a045eeb4553f5b728 Mon Sep 17 00:00:00 2001
From: Shantanu <12621235+hauntsaninja@users.noreply.github.com>
Date: Tue, 17 May 2022 00:19:44 -0700
Subject: [PATCH] feat(core): Use tomllib on Python 3.11 (#1072)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* docs: 📝 Fix typo in `pip install pdm` description (#1061)

* Use tomllib on Python 3.11

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* news

* use a compatibility module

* missed one import

Co-authored-by: t106362512 <33215526+t106362512@users.noreply.github.com>
Co-authored-by: hauntsaninja <>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 news/1072.feature.md        |  1 +
 pdm/builders/base.py        |  4 ++--
 pdm/cli/completions/pdm.ps1 | 10 ++++++++--
 pdm/cli/completions/pdm.zsh |  9 +++++++--
 pdm/compat.py               |  8 ++++++++
 pdm/formats/flit.py         |  9 ++++-----
 pdm/formats/legacy.py       | 11 ++++++-----
 pdm/formats/pipfile.py      |  4 ++--
 pdm/formats/poetry.py       |  9 ++++-----
 pdm/formats/requirements.py |  4 ++--
 pyproject.toml              |  2 +-
 11 files changed, 45 insertions(+), 26 deletions(-)
 create mode 100644 news/1072.feature.md
 create mode 100644 pdm/compat.py

diff --git a/news/1072.feature.md b/news/1072.feature.md
new file mode 100644
index 0000000000..b0bd78dda4
--- /dev/null
+++ b/news/1072.feature.md
@@ -0,0 +1 @@
+Use tomllib on Python 3.11
diff --git a/pdm/builders/base.py b/pdm/builders/base.py
index 4f9d2ffafe..e584abd3e6 100644
--- a/pdm/builders/base.py
+++ b/pdm/builders/base.py
@@ -11,9 +11,9 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterable, Mapping
 
-import tomli
 from pep517.wrappers import Pep517HookCaller
 
+from pdm.compat import tomllib
 from pdm.exceptions import BuildError
 from pdm.models.in_process import get_sys_config_paths
 from pdm.models.requirements import parse_requirement
@@ -174,7 +174,7 @@ def __init__(self, src_dir: str | Path, environment: Environment) -> None:
         logger.debug("Preparing isolated env for PEP 517 build...")
         try:
             with open(os.path.join(src_dir, "pyproject.toml"), "rb") as f:
-                spec = tomli.load(f)
+                spec = tomllib.load(f)
         except FileNotFoundError:
             spec = {}
         except Exception as e:
diff --git a/pdm/cli/completions/pdm.ps1 b/pdm/cli/completions/pdm.ps1
index db98405943..3f0759020c 100644
--- a/pdm/cli/completions/pdm.ps1
+++ b/pdm/cli/completions/pdm.ps1
@@ -125,13 +125,19 @@ function getPyPIPackages() {
 }
 
 function getPdmPackages() {
-    & $PDM_PYTHON -c "import os, re, tomli
+    & $PDM_PYTHON -c "
+import sys
+if sys.version_info >= (3, 11):
+  import tomllib
+else:
+  import tomli as tomllib
+import os, re
 PACKAGE_REGEX = re.compile(r'^[A-Za-z][A-Za-z0-9._-]*')
 def get_packages(lines):
     return [PACKAGE_REGEX.match(line).group() for line in lines]
 
 with open('pyproject.toml', 'rb') as f:
-    data = tomli.load(f)
+    data = tomllib.load(f)
 packages = get_packages(data.get('project', {}).get('dependencies', []))
 for reqs in data.get('project', {}).get('optional-dependencies', {}).values():
     packages.extend(get_packages(reqs))
diff --git a/pdm/cli/completions/pdm.zsh b/pdm/cli/completions/pdm.zsh
index 4f0787bf1a..0eb7f7def1 100644
--- a/pdm/cli/completions/pdm.zsh
+++ b/pdm/cli/completions/pdm.zsh
@@ -363,13 +363,18 @@ _pdm_groups() {
 
 _get_packages_with_python() {
   command ${PDM_PYTHON} - << EOF
-import os, re, tomli
+import sys
+if sys.version_info >= (3, 11):
+  import tomllib
+else:
+  import tomli as tomllib
+import os, re
 PACKAGE_REGEX = re.compile(r'^[A-Za-z][A-Za-z0-9._-]*')
 def get_packages(lines):
     return [PACKAGE_REGEX.match(line).group() for line in lines]
 
 with open('pyproject.toml', 'rb') as f:
-    data = tomli.load(f)
+    data = tomllib.load(f)
 packages = get_packages(data.get('project', {}).get('dependencies', []))
 for reqs in data.get('project', {}).get('optional-dependencies', {}).values():
     packages.extend(get_packages(reqs))
diff --git a/pdm/compat.py b/pdm/compat.py
new file mode 100644
index 0000000000..874c07469c
--- /dev/null
+++ b/pdm/compat.py
@@ -0,0 +1,8 @@
+import sys
+
+if sys.version_info >= (3, 11):
+    import tomllib as tomllib
+else:
+    import tomli
+
+    tomllib = tomli
diff --git a/pdm/formats/flit.py b/pdm/formats/flit.py
index 5162487646..e849635fb7 100644
--- a/pdm/formats/flit.py
+++ b/pdm/formats/flit.py
@@ -5,8 +5,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, Tuple, cast
 
-import tomli
-
+from pdm.compat import tomllib
 from pdm.formats.base import (
     MetaConverter,
     Unset,
@@ -22,8 +21,8 @@
 def check_fingerprint(project: Optional[Project], filename: PathLike) -> bool:
     with open(filename, "rb") as fp:
         try:
-            data = tomli.load(fp)
-        except tomli.TOMLDecodeError:
+            data = tomllib.load(fp)
+        except tomllib.TOMLDecodeError:
             return False
 
     return "tool" in data and "flit" in data["tool"]
@@ -145,7 +144,7 @@ def convert(
 ) -> Tuple[Mapping, Mapping]:
     with open(filename, "rb") as fp, cd(os.path.dirname(os.path.abspath(filename))):
         converter = FlitMetaConverter(
-            tomli.load(fp)["tool"]["flit"], project.core.ui if project else None
+            tomllib.load(fp)["tool"]["flit"], project.core.ui if project else None
         )
         return converter.convert()
 
diff --git a/pdm/formats/legacy.py b/pdm/formats/legacy.py
index 12af3c8241..beb76e7ab2 100644
--- a/pdm/formats/legacy.py
+++ b/pdm/formats/legacy.py
@@ -4,9 +4,8 @@
 from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, Set, Tuple, Union, cast
 
-import tomli
-
 from pdm._types import RequirementDict, Source
+from pdm.compat import tomllib
 from pdm.formats.base import (
     MetaConverter,
     Unset,
@@ -22,8 +21,8 @@
 def check_fingerprint(project: Project, filename: PathLike) -> bool:
     with open(filename, "rb") as fp:
         try:
-            data = tomli.load(fp)
-        except tomli.TOMLDecodeError:
+            data = tomllib.load(fp)
+        except tomllib.TOMLDecodeError:
             return False
 
     return (
@@ -168,7 +167,9 @@ def convert(
     project: Project, filename: Path, options: Optional[Namespace]
 ) -> Tuple[Mapping[str, Any], Mapping[str, Any]]:
     with open(filename, "rb") as fp:
-        converter = LegacyMetaConverter(tomli.load(fp)["tool"]["pdm"], project.core.ui)
+        converter = LegacyMetaConverter(
+            tomllib.load(fp)["tool"]["pdm"], project.core.ui
+        )
         return converter.convert()
 
 
diff --git a/pdm/formats/pipfile.py b/pdm/formats/pipfile.py
index 393eeb6cbe..b7155b7259 100644
--- a/pdm/formats/pipfile.py
+++ b/pdm/formats/pipfile.py
@@ -7,10 +7,10 @@
 from os import PathLike
 from typing import Any
 
-import tomli
 from packaging.markers import default_environment
 
 from pdm._types import RequirementDict
+from pdm.compat import tomllib
 from pdm.formats.base import make_array
 from pdm.models.markers import Marker
 from pdm.models.requirements import Requirement
@@ -45,7 +45,7 @@ def convert(
     project: Project, filename: PathLike, options: Namespace | None
 ) -> tuple[dict[str, Any], dict[str, Any]]:
     with open(filename, "rb") as fp:
-        data = tomli.load(fp)
+        data = tomllib.load(fp)
     result = {}
     settings = {}
     if "pipenv" in data:
diff --git a/pdm/formats/poetry.py b/pdm/formats/poetry.py
index 09f3a14eaf..bf439f0985 100644
--- a/pdm/formats/poetry.py
+++ b/pdm/formats/poetry.py
@@ -8,9 +8,8 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Mapping
 
-import tomli
-
 from pdm._types import RequirementDict, Source
+from pdm.compat import tomllib
 from pdm.formats.base import (
     MetaConverter,
     Unset,
@@ -32,8 +31,8 @@
 def check_fingerprint(project: Project | None, filename: Path | str) -> bool:
     with open(filename, "rb") as fp:
         try:
-            data = tomli.load(fp)
-        except tomli.TOMLDecodeError:
+            data = tomllib.load(fp)
+        except tomllib.TOMLDecodeError:
             return False
 
     return "tool" in data and "poetry" in data["tool"]
@@ -202,7 +201,7 @@ def convert(
 ) -> tuple[Mapping[str, Any], Mapping[str, Any]]:
     with open(filename, "rb") as fp, cd(os.path.dirname(os.path.abspath(filename))):
         converter = PoetryMetaConverter(
-            tomli.load(fp)["tool"]["poetry"], project.core.ui if project else None
+            tomllib.load(fp)["tool"]["poetry"], project.core.ui if project else None
         )
         return converter.convert()
 
diff --git a/pdm/formats/requirements.py b/pdm/formats/requirements.py
index 2fc9074722..f2aefb5ce9 100644
--- a/pdm/formats/requirements.py
+++ b/pdm/formats/requirements.py
@@ -77,11 +77,11 @@ def parse_requirement_file(
 
 
 def check_fingerprint(project: Project, filename: PathLike) -> bool:
-    import tomli
+    from pdm.compat import tomllib
 
     with open(filename, "rb") as fp:
         try:
-            tomli.load(fp)
+            tomllib.load(fp)
         except ValueError:
             # the file should be a requirements.txt if it not a TOML document.
             return True
diff --git a/pyproject.toml b/pyproject.toml
index 1a2c84cb1f..71e8846c39 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,7 +21,7 @@ dependencies = [
     "python-dotenv>=0.15",
     "resolvelib>=0.8,<0.9",
     "shellingham>=1.3.2",
-    "tomli>=1.1.0",
+    "tomli>=1.1.0; python_version < \"3.11\"",
     "tomlkit>=0.8.0,<1",
     "typing-extensions; python_version < \"3.8\"",
     "wheel>=0.36.2",