Add support for TOML files to the CLI. (#277)

* Add support for TOML files to the CLI. --------- Co-authored-by: Mahmoud Hashemi <[email protected]>
mahmoud · Nov 25, 2023 · 9082ab6 · 9082ab6
1 parent 573cb91
commit 9082ab6
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 6 deletions.
diff --git a/docs/cli.rst b/docs/cli.rst
@@ -22,8 +22,9 @@ All the power of ``glom``, without even opening your text editor!
 
    --help / -h                     show this help message and exit
    --target-file TARGET_FILE       path to target data source (optional)
-   --target-format TARGET_FORMAT   format of the source data (json or python)
-                                     (defaults to 'json')
+   --target-format TARGET_FORMAT
+                                   format of the source data (json, python, toml,
+                                   or yaml) (defaults to 'json')
    --spec-file SPEC_FILE           path to glom spec definition (optional)
    --spec-format SPEC_FORMAT       format of the glom spec definition (json, python,
                                      python-full) (defaults to 'python')

diff --git a/glom/cli.py b/glom/cli.py
@@ -11,8 +11,8 @@
   --help / -h                 show this help message and exit
   --target-file TARGET_FILE   path to target data source (optional)
   --target-format TARGET_FORMAT
-                              format of the source data (json or python)
-                              (defaults to 'json')
+                              format of the source data (json, python, toml,
+                              or yaml) (defaults to 'json')
   --spec-file SPEC_FILE       path to glom spec definition (optional)
   --spec-format SPEC_FORMAT   format of the glom spec definition (json, python,
                               python-full) (defaults to 'python')
@@ -110,8 +110,8 @@ def mw_handle_target(target_text, target_format):
     """ Handles reading in a file specified in cli command.
 
     Args:
-        target_text (str): String that specifies where 6
-        target_format (str): Valid formats include `.json` and `.yml` or `.yaml`
+        target_text (str): The target data to load, as text
+        target_format (str): Valid formats include `.json`, `.toml`, and `.yml` or `.yaml`
     Returns:
         The content of the file that you specified
     Raises:
@@ -128,6 +128,17 @@ def mw_handle_target(target_text, target_format):
             load_func = yaml.safe_load
         except ImportError:
             raise UsageError('No YAML package found. To process yaml files, run: pip install PyYAML')
+    elif target_format == 'toml':
+        missing =  UsageError('No TOML package found. To process toml files, upgrade to Python 3.11 or run: pip install tomli')
+        try:
+            import tomllib
+            load_func = tomllib.loads
+        except ImportError:
+            try:
+                import tomli
+                load_func = tomli.loads
+            except ImportError:
+                raise missing
     elif target_format == 'python':
         load_func = ast.literal_eval
     else:

diff --git a/glom/test/data/test_invalid.toml b/glom/test/data/test_invalid.toml
@@ -0,0 +1,2 @@
+# invalid
+toml = {
diff --git a/glom/test/data/test_valid.toml b/glom/test/data/test_valid.toml
@@ -0,0 +1 @@
+Hello = ["World"]
diff --git a/glom/test/test_cli.py b/glom/test/test_cli.py
@@ -78,6 +78,10 @@ def test_usage_errors(cc, basic_spec_path, basic_target_path):
     res = cc.fail_1(['glom', '--target-format', 'yaml', BASIC_SPEC, '{' + BASIC_TARGET])
     assert 'could not load target data' in res.stdout  # TODO: stderr
 
+    # bad target toml
+    res = cc.fail_1(['glom', '--target-format', 'toml', BASIC_SPEC, '{' + BASIC_TARGET])
+    assert 'could not load target data' in res.stdout  # TODO: stderr
+
     # TODO: bad target python?
 
     # bad target format  TODO: fail_2
@@ -132,6 +136,23 @@ def test_main_yaml_target():
     assert 'expected <block end>, but found' in str(excinfo.value)
 
 
+def test_main_toml_target():
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    # Handles the filepath if running tox
+    if '.tox' in cwd:
+        cwd = os.path.join(cwd.split('.tox')[0] + '/glom/test/')
+    path = os.path.join(cwd, 'data/test_valid.toml')
+    argv = ['__', '--target-file', path, '--target-format', 'toml', 'Hello']
+    assert cli.main(argv) == 0
+
+    path = os.path.join(cwd, 'data/test_invalid.toml')
+    argv = ['__', '--target-file', path, '--target-format', 'toml', 'Hello']
+    # Makes sure correct improper toml exception is raised
+    with pytest.raises(CommandLineError) as excinfo:
+        cli.main(argv)
+    assert 'Invalid initial character for a key part' in str(excinfo.value)
+
+
 def test_main_python_full_spec_python_target():
     argv = ['__', '--target-format', 'python', '--spec-format', 'python-full', 'T[T[3].bit_length()]', '{1: 2, 2: 3, 3: 4}']
     assert cli.main(argv) == 0

diff --git a/setup.py b/setup.py
@@ -40,6 +40,7 @@ def import_path(module_name, path):
       packages=['glom', 'glom.test'],
       install_requires=['boltons>=19.3.0', 'attrs', 'face==20.1.1'],
       extras_require={
+          'toml': ['tomli; python_version<"3.11"'],
           'yaml': ['PyYAML'],
       },
       entry_points={'console_scripts': ['glom = glom.cli:console_main']},