pybamm-team#1330 extend print_citation features

brosaplanella · Jan 23, 2021 · bfb93f8 · bfb93f8
1 parent bcdd19a
commit bfb93f8
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 7 deletions.
diff --git a/pybamm/citations.py b/pybamm/citations.py
@@ -5,6 +5,7 @@
 #
 import pybamm
 import os
+import pybtex
 
 
 class Citations:
@@ -74,7 +75,7 @@ def register(self, key):
             raise KeyError("'{}' is not a known citation".format(key))
         self._papers_to_cite.add(key)
 
-    def print(self, filename=None):
+    def print(self, filename=None, output_format="text"):
         """Print all citations that were used for running simulations.
 
         Parameters
@@ -84,18 +85,34 @@ def print(self, filename=None):
             terminal.
         """
         citations = ""
-        for key in self._papers_to_cite:
-            citations += self._all_citations[key] + "\n"
+        citations_file = os.path.join(pybamm.root_dir(), "pybamm", "CITATIONS.txt")
+        if output_format == "text":
+            citations = pybtex.format_from_file(
+                citations_file,
+                "plain",
+                citations=self._papers_to_cite,
+                output_backend="plaintext",
+            )
+        elif output_format == "bibtex":
+            for key in self._papers_to_cite:
+                citations += self._all_citations[key] + "\n"
+        else:
+            raise pybamm.OptionError(
+                "Output format {} not recognised. It should be 'text' or 'bibtex'.".format(
+                    output_format
+                )
+            )
+
         if filename is None:
             print(citations)
         else:
             with open(filename, "w") as f:
                 f.write(citations)
 
 
-def print_citations(filename=None):
+def print_citations(filename=None, output_format="text"):
     "See :meth:`Citations.print`"
-    pybamm.citations.print(filename)
+    pybamm.citations.print(filename, output_format)
 
 
 citations = Citations()
diff --git a/pybamm/models/submodels/electrolyte_conductivity/integrated_conductivity.py b/pybamm/models/submodels/electrolyte_conductivity/integrated_conductivity.py
@@ -30,8 +30,8 @@ class Integrated(BaseElectrolyteConductivity):
     """
 
     def __init__(self, param, domain=None):
-        pybamm.citations.register("BrosaPlanella2020")
         super().__init__(param, domain)
+        pybamm.citations.register("BrosaPlanella2020")
 
     def _higher_order_macinnes_function(self, x):
         return pybamm.log(x)

diff --git a/tests/unit/test_citations.py b/tests/unit/test_citations.py
@@ -23,9 +23,12 @@ def test_citations(self):
 
     def test_print_citations(self):
         pybamm.citations._reset()
-        pybamm.print_citations("test_citations.txt")
+        pybamm.print_citations("test_citations.txt", "text")
+        pybamm.print_citations("test_citations.txt", "bibtex")
         pybamm.citations._papers_to_cite = set()
         pybamm.print_citations()
+        with self.assertRaisesRegex(pybamm.OptionError, "'text' or 'bibtex'"):
+            pybamm.print_citations("test_citations.txt", "bad format")
 
     def test_andersson_2019(self):
         citations = pybamm.citations