Fixed python 3.12 compatibility issues. (#1004)

* correctly handle jax version == "latest"

* Changed playwright to only be a dependency if python version is less than 3.12 due to an incompatibility in its dependencies

* New pycodestyle issues with 3.12.
unittest.TestCase.assertEquals -> assertEqual

.github/workflows/dymos_docs_workflow.yml

@@ -46,7 +46,7 @@ jobs:
             SNOPT: 7.7
             OPENMDAO: 'dev'
             OPTIONAL: '[docs]'
-            JAX: '0.3.24'
+            JAX: 'latest'
             PUBLISH_DOCS: 0
 
     steps:
@@ -101,8 +101,11 @@ jobs:
           echo "============================================================="
           echo "Install jax"
           echo "============================================================="
-          python -m pip install jaxlib==${{ matrix.JAX }} jax==${{ matrix.JAX }}
-
+          if [[ "${{ matrix.JAX }}" == "latest" ]]; then
+            python -m pip install jaxlib jax
+          else
+            python -m pip install jaxlib==${{ matrix.JAX }} jax==${{ matrix.JAX }}
+          fi
       - name: Install PETSc
         if: matrix.PETSc
         shell: bash -l {0}

.github/workflows/dymos_tests_workflow.yml

@@ -76,7 +76,7 @@ jobs:
             SNOPT: 7.7
             OPENMDAO: 'dev'
             OPTIONAL: '[test]'
-            JAX: '0.3.24'
+            JAX: 'latest'
 
           # oldest supported versions
           - NAME: oldest
@@ -160,7 +160,20 @@ jobs:
           echo "============================================================="
           echo "Install jax"
           echo "============================================================="
-          python -m pip install jaxlib==${{ matrix.JAX }} jax==${{ matrix.JAX }}
+          if [[ "${{ matrix.JAX }}" == "latest" ]]; then
+            python -m pip install jaxlib jax
+          else
+            python -m pip install jaxlib==${{ matrix.JAX }} jax==${{ matrix.JAX }}
+          fi
+
+      - name: Install greenlet
+        if: env.NAME == 'latest'
+        shell: bash -l {0}
+        run: |
+          echo "============================================================="
+          echo "Install greenlet from wheels"
+          echo "============================================================="
+          pip install --only-binary :all: greenlet
 
       - name: Install PETSc
         if: env.RUN_BUILD && matrix.PETSc

dymos/trajectory/test/test_t_initial_bounds.py

@@ -146,7 +146,7 @@ def test_pair_fixed_t_initial_below(self):
 
         msg = ("'traj' <class Trajectory>: Fixed t_initial of 5.0 is outside of allowed bounds "
                "(10.0, 15.0) for phase 'phase1'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_pair_fixed_t_initial_above(self):
         kwargs = {
@@ -172,7 +172,7 @@ def test_pair_fixed_t_initial_above(self):
 
         msg = ("'traj' <class Trajectory>: Fixed t_initial of 99.0 is outside of allowed bounds "
                "(10.0, 15.0) for phase 'phase1'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_pair_t_initial_bounds_below(self):
         kwargs = {
@@ -194,7 +194,7 @@ def test_pair_t_initial_bounds_below(self):
 
         msg = ("'traj' <class Trajectory>: t_initial bounds of (5.0, 7.0) do not overlap with "
                "allowed bounds (8.0, 17.0) for phase 'phase1'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_pair_t_initial_bounds_above(self):
         kwargs = {
@@ -216,7 +216,7 @@ def test_pair_t_initial_bounds_above(self):
 
         msg = ("'traj' <class Trajectory>: t_initial bounds of (20.0, 22.0) do not overlap with "
                "allowed bounds (8.0, 17.0) for phase 'phase1'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_pair_no_duration_bounds(self):
         kwargs = {
@@ -255,7 +255,7 @@ def test_all_fixed_t_initial(self):
 
         msg = ("'traj' <class Trajectory>: Fixed t_initial of 5.0 is outside of allowed "
                "bounds (10.0, 15.0) for phase 'phase2'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_all_t_initial_bounds(self):
         nphases = 3
@@ -278,7 +278,7 @@ def test_all_t_initial_bounds(self):
 
         msg = ("'traj' <class Trajectory>: t_initial bounds of (99.0, 104.0) do not overlap with "
                "allowed bounds (10.0, 20.0) for phase 'phase2'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_odd_fixed_t_initial(self):
         nphases = 4
@@ -424,7 +424,7 @@ def test_branching_all_fixed_t_initial(self):
                "(15.0, 20.0) for phase 'br0_phase0'.\n"
                "Fixed t_initial of 10.0 is outside of allowed bounds (15.0, 20.0) for phase "
                "'br1_phase0'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_branching_all_t_initial_bounds(self):
         nphases = 3  # number of phases in trunk and each branch
@@ -473,7 +473,7 @@ def test_branching_all_t_initial_bounds(self):
                "allowed bounds (20.0, 30.0) for phase 'br0_phase1'.\n"
                "t_initial bounds of (0.0, 5) do not overlap with allowed bounds (20.0, 30.0) "
                "for phase 'br1_phase1'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)
 
     def test_branching_odd_fixed_t_initial(self):
         nphases = 3  # number of phases in trunk and each branch
@@ -521,4 +521,4 @@ def test_branching_odd_fixed_t_initial(self):
                "(40.0, 50.0) for phase 'br0_phase3'.\n"
                "Fixed t_initial of 60.0 is outside of allowed bounds (40.0, 50.0) for phase "
                "'br1_phase3'.")
-        self.assertEquals(cm.exception.args[0], msg)
+        self.assertEqual(cm.exception.args[0], msg)

dymos/transcriptions/pseudospectral/radau_pseudospectral.py

@@ -253,7 +253,7 @@ def _get_rate_source_path(self, state_name, nodes, phase):
         try:
             var = phase.state_options[state_name]['rate_source']
         except RuntimeError:
-            raise ValueError(f"state '{state_name}' in phase '{ phase.name}' was not given a "
+            raise ValueError(f"state '{state_name}' in phase '{phase.name}' was not given a "
                              "rate_source")
 
         # Note the rate source must be shape-compatible with the state

dymos/utils/test/test_testing_utils.py

@@ -73,8 +73,8 @@ def test_unequal_time_series_rel_only(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with {start_of_expected_errmsg} but "
                         f"instead was {actual_errmsg}")
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 2)
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 0)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 0)
 
     def test_unequal_time_series_abs_only(self):
         # slightly modify the "to be checked" time series and check that the assert is working
@@ -102,8 +102,8 @@ def test_unequal_time_series_abs_only(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with {start_of_expected_errmsg} but "
                         f"instead was {actual_errmsg}")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 2)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 0)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 0)
 
     def test_unequal_time_series_abs_and_rel(self):
         # slightly modify the "to be checked" time series and check that the assert is working
@@ -139,8 +139,8 @@ def test_unequal_time_series_abs_and_rel(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with '{start_of_expected_errmsg}' but "
                         f"instead was '{actual_errmsg}'")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 2)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 2)
 
         # for > 100, uses the rel, x_check[15] is ~ 150
         x_check[5] = x_check_5_orig
@@ -166,8 +166,8 @@ def test_unequal_time_series_abs_and_rel(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with '{start_of_expected_errmsg}' but "
                         f"instead was '{actual_errmsg}'")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 2)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 2)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 2)
 
         # Combine the two cases where one data paint fails because of abs error and one because
         #   of rel error
@@ -187,8 +187,8 @@ def test_unequal_time_series_abs_and_rel(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with '{start_of_expected_errmsg}' but "
                         f"instead was '{actual_errmsg}'")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 3)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 3)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 3)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 3)
 
     def test_no_overlapping_time(self):
         t_ref, x_ref = create_linear_time_series(100, 0.0, 500.0, 0.0, 1000.0)
@@ -253,8 +253,8 @@ def test_multi_dimensional_unequal(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with {start_of_expected_errmsg} but "
                         f"instead was {actual_errmsg}")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 0)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 199)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 0)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 199)
 
     def test_multi_dimensional_unequal_abs_and_rel(self):
         t_ref, x_ref_1 = create_linear_time_series(10, 0.0, 500.0, 0.0, 1000.0)
@@ -279,8 +279,8 @@ def test_multi_dimensional_unequal_abs_and_rel(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with {start_of_expected_errmsg} but "
                         f"instead was {actual_errmsg}")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 19)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 19)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 19)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 19)
 
     def test_multi_dimensional_with_overlapping_times(self):
         t_ref, x_ref_1 = create_linear_time_series(100, 0.0, 500.0, 0.0, 1000.0)
@@ -314,8 +314,8 @@ def test_multi_dimensional_unequal_with_overlapping_times(self):
         self.assertTrue(actual_errmsg.startswith(start_of_expected_errmsg),
                         f"Error message expected to start with {start_of_expected_errmsg} but "
                         f"instead was {actual_errmsg}")
-        self.assertEquals(actual_errmsg.count('>ABS_TOL'), 0)
-        self.assertEquals(actual_errmsg.count('>REL_TOL'), 99)
+        self.assertEqual(actual_errmsg.count('>ABS_TOL'), 0)
+        self.assertEqual(actual_errmsg.count('>REL_TOL'), 99)
 
 
 @use_tempdirs

dymos/visualization/linkage/test/test_gui.py

@@ -1,9 +1,16 @@
 """Test Dymos Linkage report GUI with using Playwright."""
 import unittest
 import os
+try:
+    import playwright
+except ImportError:
+    playwright = None
+
+
+if playwright is not None:
+    os.system("playwright install")
+    from linkage_report_ui_test import dymos_linkage_gui_test_case  # nopep8: E402
 
-os.system("playwright install")
-from linkage_report_ui_test import dymos_linkage_gui_test_case  # nopep8: E402
 
 if __name__ == "__main__":
     unittest.main()

dymos/visualization/timeseries_plots.py

@@ -136,7 +136,7 @@ def _mpl_timeseries_plots(time_units, var_units, phase_names, phases_node_path,
         plt.subplots_adjust(bottom=0.23, top=0.9, left=0.2)
 
         # save to file
-        plot_file_path = plot_dir_path.joinpath(f'{var_name.replace(":","_")}.png')
+        plot_file_path = plot_dir_path.joinpath(f'{var_name.replace(":", "_")}.png')
         plt.savefig(plot_file_path, dpi=dpi)
         plt.close(fig)
         plotfiles.append(plot_file_path)

setup.py

@@ -1,3 +1,5 @@
+from packaging.version import Version
+from platform import python_version
 from setuptools import find_packages, setup
 
 # Setup optional dependencies
@@ -27,11 +29,13 @@
         'testflo>=1.3.6',
         'matplotlib',
         'numpydoc>=1.1',
-        'playwright>=1.20',
-        'aiounittest'
     ]
 }
 
+# playwright dependencies are currently incompatible with python 3.12
+if Version(python_version()) < Version('3.12.0'):
+    optional_dependencies['test'].extend(['playwright>=1.20', 'aiounittest'])
+
 # Add an optional dependency that concatenates all others
 optional_dependencies['all'] = sorted([
     dependency