From 168d3bdd3c038529c6eca0ef58ab6f6eb2216e52 Mon Sep 17 00:00:00 2001
From: Axel Kohlmeyer <akohlmey@gmail.com>
Date: Sat, 14 May 2022 10:13:08 -0400
Subject: [PATCH] Correctly handle the case that the YAML python module is not
 installed

---
 python/lammps/formats.py          | 17 +++++++++++++----
 unittest/python/python-formats.py | 16 +++++++++++-----
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/python/lammps/formats.py b/python/lammps/formats.py
index d34998f01c..3377bccb41 100644
--- a/python/lammps/formats.py
+++ b/python/lammps/formats.py
@@ -17,11 +17,18 @@
 # and Axel Kohlmeyer <akohlmey@gmail.com>
 ################################################################################
 
-import re, yaml
+import re
+
+has_yaml = False
 try:
-  from yaml import CSafeLoader as Loader
-except ImportError:
-  from yaml import SafeLoader as Loader
+  import yaml
+  has_yaml = True
+  try:
+    from yaml import CSafeLoader as Loader
+  except ImportError:
+    from yaml import SafeLoader as Loader
+except ImportError: 
+  pass
 
 class LogFile:
   """Reads LAMMPS log files and extracts the thermo information
@@ -66,6 +73,8 @@ class LogFile:
                 style = LogFile.STYLE_YAML
                 yamllog += line;
                 current_run = {}
+                if not has_yaml:
+                  raise Exception('Cannot process YAML format logs without the PyYAML Python module')
 
             elif re.match(r'^\.\.\.$', line):
                 thermo = yaml.load(yamllog, Loader=Loader)
diff --git a/unittest/python/python-formats.py b/unittest/python/python-formats.py
index c3dc6d52dc..b34f6db36d 100644
--- a/unittest/python/python-formats.py
+++ b/unittest/python/python-formats.py
@@ -2,11 +2,16 @@ import os
 import unittest
 from lammps.formats import LogFile, AvgChunkFile
 
-import yaml
+has_yaml = False
 try:
-    from yaml import CSafeLoader as Loader, CSafeDumper as Dumper
-except ImportError:
-    from yaml import SafeLoader, SafeDumper
+    import yaml
+    has_yaml = True
+    try:
+        from yaml import CSafeLoader as Loader, CSafeDumper as Dumper
+    except ImportError:
+        from yaml import SafeLoader, SafeDumper
+except:
+    pass
 
 EXAMPLES_DIR=os.path.abspath(os.path.join(__file__, '..', '..', '..', 'examples'))
 
@@ -65,6 +70,7 @@ class Logfiles(unittest.TestCase):
 
         self.assertEqual(run0["Step"], list(range(0,350, 50)))
 
+    @unittest.skipIf(not has_yaml,"Missing the PyYAML python module")
     def testYamlLogFile(self):
         log = LogFile(os.path.join(EXAMPLES_DIR, YAML_STYLE_EXAMPLE_LOG))
         self.assertEqual(len(log.runs), 2)
@@ -127,7 +133,7 @@ try:
 except:
     pass
 
-@unittest.skipIf(not has_dump_yaml, "Either atom_style full or dump_style yaml are not available")
+@unittest.skipIf(not (has_dump_yaml and has_yaml), "Either atom_style full, dump_style yaml, or the python PyYAML module are not available")
 class PythonDump(unittest.TestCase):
     def setUp(self):
         machine = None