From 0ee44c8abe4b91781ef5a84ce12cbad438db9236 Mon Sep 17 00:00:00 2001
From: "Richy Pitman (CCN RPT RES3)" <rpit@equinor.com>
Date: Fri, 15 Dec 2023 17:43:21 +0100
Subject: [PATCH] Refactor imports and plotter function  after failing isort
 and mypy tests

---
 pyscal/plotting.py     | 43 +++++++++++++++++++++++++++++-------------
 pyscal/pyscalcli.py    |  2 +-
 tests/test_plotting.py |  2 +-
 3 files changed, 32 insertions(+), 15 deletions(-)

diff --git a/pyscal/plotting.py b/pyscal/plotting.py
index 8be7d9d0..d42e5150 100644
--- a/pyscal/plotting.py
+++ b/pyscal/plotting.py
@@ -14,7 +14,7 @@
 import matplotlib.pyplot as plt
 import pandas as pd
 
-from pyscal import PyscalList, WaterOil, GasWater, GasOil, WaterOilGas
+from pyscal import GasOil, GasWater, PyscalList, WaterOil, WaterOilGas
 
 # Data for configuring plot based on pyscal model type
 PLOT_CONFIG_OPTIONS = {
@@ -79,9 +79,9 @@ def format_relperm_plot(fig: plt.Figure, **kwargs) -> plt.Figure:
     # permeability used to calculate kr
     if kwargs["semilog"]:
         ax.set_yscale("log")
-        ax.set_ylim([1e-6, 1])
+        ax.set_ylim((1e-6, 1.0))
     else:
-        ax.set_ylim([0, 1])
+        ax.set_ylim((0.0, 1.0))
 
     # Add legend
     plt.legend()
@@ -116,11 +116,11 @@ def format_cap_pressure_plot(
     ax.set_ylabel(kwargs["pc_name"].lower().capitalize())
 
     # Set axis limits
-    ax.set_xlim([0, 1])
+    ax.set_xlim((0.0, 1.0))
 
     # Set lower y-axis limit to 0 if Pc >= 0
     if not neg_pc:
-        ax.set_ylim(bottom=0)
+        ax.set_ylim(bottom=0.0)
 
     return fig
 
@@ -331,24 +331,41 @@ def plotter(
     kwargs = {"pc": pc, "semilog": semilog}
 
     for model in models.pyscal_list:
-        # Get SATNUM number as an integer. Used in the naming of saved figures
-        satnum = get_satnum_from_tag(model.tag)
-
         if isinstance(model, WaterOilGas):
-            plot_individual_curves("WaterOil", model.wateroil.table, satnum, **kwargs)
-            plot_individual_curves("GasOil", model.gasoil.table, satnum, **kwargs)
+            # the wateroil and gasoil instance variables are optional for the
+            # WaterOilGas class
+            if model.wateroil:
+                plot_individual_curves(
+                    "WaterOil",
+                    model.wateroil.table,
+                    get_satnum_from_tag(model.wateroil.tag),
+                    **kwargs,
+                )
+            if model.gasoil:
+                plot_individual_curves(
+                    "GasOil",
+                    model.gasoil.table,
+                    get_satnum_from_tag(model.gasoil.tag),
+                    **kwargs,
+                )
 
         elif isinstance(model, WaterOil):
-            plot_individual_curves("WaterOil", model.table, satnum, **kwargs)
+            plot_individual_curves(
+                "WaterOil", model.table, get_satnum_from_tag(model.tag), **kwargs
+            )
 
         elif isinstance(model, GasOil):
-            plot_individual_curves("WaterOil", model.table, satnum, **kwargs)
+            plot_individual_curves(
+                "WaterOil", model.table, get_satnum_from_tag(model.tag), **kwargs
+            )
 
         elif isinstance(model, GasWater):
             # The GasWater object has a different structure to the others, and
             # requires formatting
             table = format_gaswater_table(model)
-            plot_individual_curves("GasWater", table, satnum, **kwargs)
+            plot_individual_curves(
+                "GasWater", table, get_satnum_from_tag(model.wateroil.tag), **kwargs
+            )
 
         else:
             raise Exception(
diff --git a/pyscal/pyscalcli.py b/pyscal/pyscalcli.py
index 0d25caec..e34696b0 100644
--- a/pyscal/pyscalcli.py
+++ b/pyscal/pyscalcli.py
@@ -13,9 +13,9 @@
     GasWater,
     SCALrecommendation,
     WaterOilGas,
-    plotting,
     __version__,
     getLogger_pyscal,
+    plotting,
 )
 
 from .factory import PyscalFactory
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 2d5a8938..5bbd26f9 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -1,6 +1,6 @@
 import pytest
 
-from pyscal import plotting, PyscalList, WaterOil, GasWater, GasOil, WaterOilGas
+from pyscal import GasOil, GasWater, PyscalList, WaterOil, WaterOilGas, plotting
 
 
 def test_get_satnum_from_tag():