Merge pull request #916 from hschilling/P165313791-N2-title

Add user option to set the title of the N2 diagram
OpenMDAO · Apr 24, 2019 · 22e24e8 · 22e24e8
2 parents 2c1a57e + 0138b9f
commit 22e24e8
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 3 deletions.
diff --git a/openmdao/devtools/html_utils.py b/openmdao/devtools/html_utils.py
@@ -346,7 +346,11 @@ def __init__(self, filename, embeddable=False, title=None, styles=None):
                 self.template = '\n\n'.join([style_elems, template])
             else:
                 meta = '<meta http-equiv="Content-Type" content="text/html; charset=UTF-8">'
-                head = '\n\n'.join([meta, style_elems])  # Write styles to head
+                if title:
+                    title_tag = "<title>%s</title>" % title
+                    head = '\n\n'.join([title_tag, meta, style_elems])  # Write styles to head
+                else:
+                    head = '\n\n'.join([meta, style_elems])  # Write styles to head
                 self.template = head_and_body(head=head, body=template)
 
         if title is not None:

diff --git a/openmdao/devtools/problem_viewer/problem_viewer.py b/openmdao/devtools/problem_viewer/problem_viewer.py
@@ -259,7 +259,8 @@ def view_tree(*args, **kwargs):
     view_model(*args, **kwargs)
 
 
-def view_model(data_source, outfile='n2.html', show_browser=True, embeddable=False):
+def view_model(data_source, outfile='n2.html', show_browser=True, embeddable=False,
+               title=None):
     """
     Generates an HTML file containing a tree viewer.
 
@@ -280,6 +281,10 @@ def view_model(data_source, outfile='n2.html', show_browser=True, embeddable=Fal
     embeddable : bool, optional
         If True, gives a single HTML file that doesn't have the <html>, <DOCTYPE>, <body>
         and <head> tags. If False, gives a single, standalone HTML file for viewing.
+
+    title : str, optional
+        The title for the diagram. Used in the HTML title and also shown on the page.
+
     """
     # grab the model viewer data
     model_data = _get_viewer_data(data_source)
@@ -305,8 +310,13 @@ def view_model(data_source, outfile='n2.html', show_browser=True, embeddable=Fal
     with open(os.path.join(style_dir, "fontello.woff"), "rb") as f:
         encoded_font = str(base64.b64encode(f.read()).decode("ascii"))
 
+    if title:
+        title = "OpenMDAO Model Hierarchy and N2 diagram: %s" % title
+    else:
+        title = "OpenMDAO Model Hierarchy and N2 diagram"
+
     h = DiagramWriter(filename=os.path.join(vis_dir, "index.html"),
-                      title="OpenMDAO Model Hierarchy and N<sup>2</sup> diagram.",
+                      title=title,
                       styles=styles, embeddable=embeddable)
 
     # put all style and JS into index

diff --git a/openmdao/devtools/problem_viewer/tests/test_viewmodeldata.py b/openmdao/devtools/problem_viewer/tests/test_viewmodeldata.py
@@ -198,6 +198,22 @@ def test_view_model_command(self):
         filename = os.path.abspath(sellar.__file__).replace('.pyc', '.py')  # PY2
         check_call('openmdao view_model --no_browser %s' % filename)
 
+    def test_view_model_set_title(self):
+        """
+        Test that an n2 html file is generated from a Problem.
+        """
+        p = Problem()
+        p.model = SellarStateConnection()
+        p.setup(check=False)
+        view_model(p, outfile=self.problem_html_filename, show_browser=DEBUG,
+                   title="Sellar State Connection")
+
+        # Check that the html file has been created and has something in it.
+        self.assertTrue(os.path.isfile(self.problem_html_filename),
+                        (self.problem_html_filename + " is not a valid file."))
+        self.assertTrue( 'OpenMDAO Model Hierarchy and N2 diagram: Sellar State Connection' \
+                         in open(self.problem_html_filename).read() )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/openmdao/utils/om.py b/openmdao/utils/om.py
@@ -44,6 +44,8 @@ def _view_model_setup_parser(parser):
                         help="don't display in a browser.")
     parser.add_argument('--embed', action='store_true', dest='embeddable',
                         help="create embeddable version.")
+    parser.add_argument('--title', default=None,
+                        action='store', dest='title', help='diagram title.')
 
 
 def _view_model_cmd(options):
@@ -62,6 +64,7 @@ def _view_model_cmd(options):
         def _viewmod(prob):
             view_model(prob, outfile=options.outfile,
                        show_browser=not options.no_browser,
+                       title=options.title,
                        embeddable=options.embeddable)
             exit()  # could make this command line selectable later
 
@@ -71,6 +74,7 @@ def _viewmod(prob):
     else:
         # assume the file is a recording, run standalone
         view_model(filename, outfile=options.outfile,
+                   title=options.title,
                    show_browser=not options.no_browser,
                    embeddable=options.embeddable)