Merge 473782b into b94bbcc

CHANGES.rst

@@ -1,6 +1,17 @@
 Changes
 =======
 
+``v1.1``
+----------
+*released 2020-12-20*
+
+Features
+~~~~~~~~
+
+- **utils:**  Enable supplementary_imports in VariableWriter and EquationWriter
+  (`PR #84 <https://github.com/environmentalscience/essm/pull/96>`__)
+
+
 ``v1.0.1``
 ----------
 *released 2020-11-05*

essm/_generator.py

@@ -121,7 +121,8 @@ class VariableWriter(object):
     .. code-block:: python
         from essm._generator import VariableWriter
         from essm.variables.physics.thermodynamics import c_pa, P_a
-        writer = VariableWriter(docstring="Test.")
+        writer = VariableWriter(docstring="Test.",
+            supplementary_imports={'sympy.physics.units': {'mega', 'kilo'}})
         writer.var(c_pa)
         writer.var(P_a)
         print(writer)
@@ -133,11 +134,12 @@ class VariableWriter(object):
         'essm.variables': {'Variable'},
     }
 
-    def __init__(self, docstring=None):
+    def __init__(self, docstring=None, supplementary_imports={}):
         """Initialize variable writer."""
         self.docstring = docstring
         self._imports = defaultdict(set)
         self._imports.update(**self.default_imports)
+        self._imports.update(**supplementary_imports)
         self.vars = []
 
     @property
@@ -217,7 +219,8 @@ def var(self, var1):
         .. code-block:: python
             from essm._generator import VariableWriter
             from essm.variables.physics.thermodynamics import c_pa, P_a
-            writer = VariableWriter(docstring="Test.")
+            writer = VariableWriter(docstring="Test.",
+                supplementary_imports={'sympy.physics.units': {'mega'}})
             writer.var(c_pa)
             writer.var(P_a)
             print(writer)
@@ -248,12 +251,13 @@ class EquationWriter(object):
         from essm.equations import Equation
         from essm._generator import EquationWriter
         from essm.variables.units import second, meter, kelvin
-        from essm.variables.physics.thermodynamics import R_s, D_va, T_a, \
-            P_a, P_wa, P_N2, P_O2
+        from essm.variables.physics.thermodynamics import (R_s, D_va, T_a,
+            P_a, P_wa, P_N2, P_O2)
         from essm.equations.physics.thermodynamics import eq_Pa
         from sympy import Eq, symbols
         p_Dva1, p_Dva2 = symbols('p_Dva1, p_Dva2')
-        writer = EquationWriter(docstring="Test.")
+        writer = EquationWriter(docstring="Test.",
+            supplementary_imports={'sympy': {'exp', 'sin'}})
         writer.eq(eq_Pa)
         writer.neweq('eq_Pwa_Pa', Eq(P_wa, P_a - P_N2 - P_O2),
                 doc='Calculate P_wa from total air pressure.',
@@ -277,11 +281,12 @@ class EquationWriter(object):
     }
     """Set up default imports, including standard division."""
 
-    def __init__(self, docstring=None):
+    def __init__(self, docstring=None, supplementary_imports={}):
         """Initialise equation writer."""
         self.docstring = docstring
         self._imports = defaultdict(set)
         self._imports.update(**self.default_imports)
+        self._imports.update(**supplementary_imports)
         self.eqs = []
 
     @property
@@ -371,15 +376,16 @@ def neweq(self, name, expr, doc='', parents=None, variables=None):
             self._imports['essm'].add(match.group())
 
     def eq(self, eq1):
-        """Add pre-defined equation to writer, including any internal variables.
+        """Add pre-defined equation to writer, incl. any internal variables.
 
         Example:
 
         .. code-block:: python
             from essm._generator import EquationWriter, VariableWriter
             from essm.variables.utils import get_parents
             from essm.equations.leaf.energy_water import eq_Pwl, eq_Cwl
-            writer = EquationWriter(docstring="Test.")
+            writer = EquationWriter(docstring="Test.",
+                supplementary_imports={'sympy': {'exp', 'sin'}})
             writer.eq(eq_Pwl)
             writer.eq(eq_Cwl)
             print(writer)

tests/test_equations.py

@@ -392,3 +392,40 @@ def test_equation_writer_linebreaks(tmpdir):
     with open(eq_file.strpath) as outfile:
         maxlinelength = max([len(line) for line in outfile.readlines()])
     assert maxlinelength < 80
+
+
+def test_equation_variable_writer_imports(tmpdir):
+    """Provide additional imports in EquationWriter and VariableWriter"""
+    from sympy import Eq, var, sin
+    from sympy.physics.units import mega, joule
+
+    g = {}
+
+    class energy1(Variable):
+        """Energy variable"""
+        unit = mega*joule
+
+    class eq_sin_energy(Equation):
+        """Equation with sine function"""
+        expr = Eq(demo_1, sin(demo_d/demo_d1))
+
+    writer_var = VariableWriter(docstring='Test of VariableWriter.',
+                                supplementary_imports={
+                                    'sympy.physics.units': {'mega'}})
+    writer_var.var(energy1)
+    var_file = tmpdir.mkdir('test1').join('test_vars.py')
+    writer_var.write(var_file.strpath)
+    with open(var_file, "rb") as source_file:
+        code = compile(var_file.read(), var_file, "exec")
+    exec(code, g)
+    assert g['energy1'].definition.unit == mega*joule
+
+    writer_eq = EquationWriter(docstring='Test of EquationWriter.',
+                               supplementary_imports={'sympy': {'sin'}})
+    writer_eq.eq(eq_sin_energy)
+    eq_file = tmpdir.mkdir('test2').join('test_equations_sin.py')
+    writer_eq.write(eq_file.strpath)
+    with open(eq_file, "rb") as source_file:
+        code = compile(eq_file.read(), eq_file, "exec")
+    exec(code, g)
+    assert g['eq_sin_energy'].definition.expr == eq_sin_energy.definition.expr