From 5909ba87f79ea1d687e5836c62dc82a116f4229f Mon Sep 17 00:00:00 2001 From: Eric McDonald <221418+emcd@users.noreply.github.com> Date: Wed, 30 Nov 2022 14:23:50 -0800 Subject: [PATCH] Prevent crashes against valid 'pyproject.toml'. (#1040) * Replace 'toml' dependency with 'tomli', which fully supports TOML 1. Co-authored-by: Eric McDonald --- setup.py | 2 +- yapf/yapflib/file_resources.py | 15 ++++++++------- yapf/yapflib/style.py | 25 ++++++++++++++----------- yapftests/file_resources_test.py | 10 +++++----- yapftests/style_test.py | 4 ++-- 5 files changed, 30 insertions(+), 26 deletions(-) diff --git a/yapf/yapflib/file_resources.py b/yapf/yapflib/file_resources.py index b5e2612b..6809ca9f 100644 --- a/yapf/yapflib/file_resources.py +++ b/yapf/yapflib/file_resources.py @@ -49,14 +49,15 @@ def _GetExcludePatternsFromPyprojectToml(filename): """Get a list of file patterns to ignore from pyproject.toml.""" ignore_patterns = [] try: - import toml + import tomli as tomllib except ImportError: raise errors.YapfError( - "toml package is needed for using pyproject.toml as a " + "tomli package is needed for using pyproject.toml as a " "configuration file") if os.path.isfile(filename) and os.access(filename, os.R_OK): - pyproject_toml = toml.load(filename) + with open(filename, 'rb') as fd: + pyproject_toml = tomllib.load(fd) ignore_patterns = pyproject_toml.get('tool', {}).get('yapfignore', {}).get('ignore_patterns', []) @@ -127,19 +128,19 @@ def GetDefaultStyleForDir(dirname, default_style=style.DEFAULT_STYLE): # See if we have a pyproject.toml file with a '[tool.yapf]' section. config_file = os.path.join(dirname, style.PYPROJECT_TOML) try: - fd = open(config_file) + fd = open(config_file, 'rb') except IOError: pass # It's okay if it's not there. else: with fd: try: - import toml + import tomli as tomllib except ImportError: raise errors.YapfError( - "toml package is needed for using pyproject.toml as a " + "tomli package is needed for using pyproject.toml as a " "configuration file") - pyproject_toml = toml.load(config_file) + pyproject_toml = tomllib.load(fd) style_dict = pyproject_toml.get('tool', {}).get('yapf', None) if style_dict is not None: return config_file diff --git a/yapf/yapflib/style.py b/yapf/yapflib/style.py index 233a64e6..c8397b32 100644 --- a/yapf/yapflib/style.py +++ b/yapf/yapflib/style.py @@ -746,17 +746,18 @@ def _CreateConfigParserFromConfigFile(config_filename): # Provide a more meaningful error here. raise StyleConfigError( '"{0}" is not a valid style or file path'.format(config_filename)) - with open(config_filename) as style_file: - config = py3compat.ConfigParser() - if config_filename.endswith(PYPROJECT_TOML): - try: - import toml - except ImportError: - raise errors.YapfError( - "toml package is needed for using pyproject.toml as a " - "configuration file") - - pyproject_toml = toml.load(style_file) + config = py3compat.ConfigParser() + + if config_filename.endswith(PYPROJECT_TOML): + try: + import tomli as tomllib + except ImportError: + raise errors.YapfError( + "tomli package is needed for using pyproject.toml as a " + "configuration file") + + with open(config_filename, 'rb') as style_file: + pyproject_toml = tomllib.load(style_file) style_dict = pyproject_toml.get("tool", {}).get("yapf", None) if style_dict is None: raise StyleConfigError( @@ -766,7 +767,9 @@ def _CreateConfigParserFromConfigFile(config_filename): config.set('style', k, str(v)) return config + with open(config_filename) as style_file: config.read_file(style_file) + if config_filename.endswith(SETUP_CONFIG): if not config.has_section('yapf'): raise StyleConfigError( diff --git a/yapftests/file_resources_test.py b/yapftests/file_resources_test.py index 31184c4a..f54f393d 100644 --- a/yapftests/file_resources_test.py +++ b/yapftests/file_resources_test.py @@ -75,7 +75,7 @@ def test_get_exclude_file_patterns_from_yapfignore_with_wrong_syntax(self): def test_get_exclude_file_patterns_from_pyproject(self): try: - import toml + import tomli except ImportError: return local_ignore_file = os.path.join(self.test_tmpdir, 'pyproject.toml') @@ -93,7 +93,7 @@ def test_get_exclude_file_patterns_from_pyproject(self): @unittest.skipUnless(py3compat.PY36, 'Requires Python 3.6') def test_get_exclude_file_patterns_from_pyproject_with_wrong_syntax(self): try: - import toml + import tomli except ImportError: return local_ignore_file = os.path.join(self.test_tmpdir, 'pyproject.toml') @@ -109,7 +109,7 @@ def test_get_exclude_file_patterns_from_pyproject_with_wrong_syntax(self): def test_get_exclude_file_patterns_from_pyproject_no_ignore_section(self): try: - import toml + import tomli except ImportError: return local_ignore_file = os.path.join(self.test_tmpdir, 'pyproject.toml') @@ -122,7 +122,7 @@ def test_get_exclude_file_patterns_from_pyproject_no_ignore_section(self): def test_get_exclude_file_patterns_from_pyproject_ignore_section_empty(self): try: - import toml + import tomli except ImportError: return local_ignore_file = os.path.join(self.test_tmpdir, 'pyproject.toml') @@ -192,7 +192,7 @@ def test_setup_config(self): def test_pyproject_toml(self): # An empty pyproject.toml file should not be used try: - import toml + import tomli except ImportError: return diff --git a/yapftests/style_test.py b/yapftests/style_test.py index 8a37f953..d2203d9a 100644 --- a/yapftests/style_test.py +++ b/yapftests/style_test.py @@ -230,7 +230,7 @@ def testErrorUnknownStyleOption(self): def testPyprojectTomlNoYapfSection(self): try: - import toml + import tomli except ImportError: return @@ -242,7 +242,7 @@ def testPyprojectTomlNoYapfSection(self): def testPyprojectTomlParseYapfSection(self): try: - import toml + import tomli except ImportError: return