use single Kaleido instance for making baselines (faster performance)

Author: emilykl

test/image/make_baseline.py

@@ -1,9 +1,14 @@
+import asyncio
 import os
 import sys
 import json
+
 import plotly.io as pio
+import kaleido
+
 from convert_b64 import arraysToB64
 
+
 args = []
 if len(sys.argv) == 2 :
     args = sys.argv[1].split()
@@ -31,11 +36,12 @@
 print('output to', dirOut)
 
 mathjax_version = 2
+mathjax = None
 if 'mathjax3' in sys.argv or 'mathjax3=' in sys.argv :
     # until https://github.com/plotly/Kaleido/issues/124 is addressed
     # we are uanble to use local mathjax v3 installed in node_modules
     # for now let's download it from the internet:
-    pio.defaults.mathjax = 'https://cdn.jsdelivr.net/npm/mathjax@3.2.2/es5/tex-svg.js'
+    mathjax = 'https://cdn.jsdelivr.net/npm/mathjax@3.2.2/es5/tex-svg.js'
     mathjax_version = 3
     print('Kaleido using MathJax v3')
 
@@ -52,8 +58,6 @@
 
     plotlyjs = plotlyjs_with_virtual_webgl
 
-pio.defaults.plotlyjs = plotlyjs
-
 pio.templates.default = 'none'
 
 ALL_MOCKS = [os.path.splitext(a)[0] for a in os.listdir(dirIn) if a.endswith('.json')]
@@ -71,6 +75,7 @@
     'map_predefined-styles2',  # Temporarily blacklist: fails with Kaleido v1.0.0rc14
     'grid_subplot_types', # Temporarily blacklist: fails with Kaleido v1.0.0rc14
     'map_fonts-supported-metropolis', # Temporarily blacklist: fails with Kaleido v1.0.0rc14
+    'map_fonts-supported-metropolis-italic', # Temporarily blacklist: fails with Kaleido v1.0.0rc14
     'map_fonts-supported-metropolis-weight', # Temporarily blacklist: fails with Kaleido v1.0.0rc14
     'map_fonts-supported-open-sans-weight', # Temporarily blacklist: fails with Kaleido v1.0.0rc14
 ]
@@ -81,55 +86,74 @@
     sys.exit(1)
 
 failed = []
-for name in allNames :
-    outName = name
-    if mathjax_version == 3 :
-        outName = 'mathjax3___' + name
-
-    print(outName)
-
-    created = False
-
-    MAX_RETRY = 2 # 1 means retry once
-    for attempt in range(0, MAX_RETRY + 1) :
-        with open(os.path.join(dirIn, name + '.json'), 'r') as _in :
-            fig = json.load(_in)
-
-            width = 700
-            height = 500
-            if 'layout' in fig :
-                layout = fig['layout']
-                if 'autosize' not in layout or layout['autosize'] != True :
-                    if 'width' in layout :
-                        width = layout['width']
-                    if 'height' in layout :
-                        height = layout['height']
-
-            if 'b64' in sys.argv or 'b64=' in sys.argv or 'b64-json' in sys.argv :
-                newFig = dict()
-                arraysToB64(fig, newFig)
-                fig = newFig
-                if 'b64-json' in sys.argv and attempt == 0 : print(json.dumps(fig, indent = 2))
-
-            try :
-                pio.write_image(
-                    fig=fig,
-                    file=os.path.join(dirOut, outName + '.png'),
-                    width=width,
-                    height=height,
-                    validate=False
-                )
-                created = True
-            except Exception as e :
-                print(e)
-                if attempt < MAX_RETRY :
-                    print('retry', attempt + 1, '/', MAX_RETRY)
-                else :
-                    failed.append(outName)
-
-        if(created) : break
-
-if len(failed) > 0 :
-    print('Failed at :')
-    print(failed)
-    sys.exit(1)
+
+async def make_baselines_async():
+
+    kopts = dict(
+        plotlyjs=plotlyjs,
+    )
+    if mathjax is not None:
+        kopts['mathjax'] = mathjax
+
+    async with kaleido.Kaleido(n=1, **kopts) as k:
+        for name in allNames:
+            outName = name
+            if mathjax_version == 3:
+                outName = 'mathjax3___' + name
+
+            print(outName)
+
+            created = False
+
+            MAX_RETRY = 2 # 1 means retry once
+            for attempt in range(0, MAX_RETRY + 1) :
+                with open(os.path.join(dirIn, name + '.json'), 'r') as _in :
+                    fig = json.load(_in)
+
+                    width = 700
+                    height = 500
+                    if 'layout' in fig :
+                        layout = fig['layout']
+                        if 'autosize' not in layout or layout['autosize'] != True :
+                            if 'width' in layout :
+                                width = layout['width']
+                            if 'height' in layout :
+                                height = layout['height']
+
+                    if 'b64' in sys.argv or 'b64=' in sys.argv or 'b64-json' in sys.argv :
+                        newFig = dict()
+                        arraysToB64(fig, newFig)
+                        fig = newFig
+                        if 'b64-json' in sys.argv and attempt == 0 : print(json.dumps(fig, indent = 2))
+
+                    try:
+                        bytes = await k.calc_fig(
+                            fig,
+                            path=None,
+                            opts=dict(
+                                format="png",
+                                width=width,
+                                height=height,
+                            ),
+                        )
+                        filename = os.path.join(dirOut, outName + '.png')
+                        with open(filename, "wb") as f:
+                            f.write(bytes)
+                        created = True
+                    except Exception as e:
+                        print(e)
+                        if attempt < MAX_RETRY :
+                            print('retry', attempt + 1, '/', MAX_RETRY)
+                        else :
+                            failed.append(outName)
+
+                if(created): break
+
+    if len(failed) > 0 :
+        print('Failed at :')
+        print(failed)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    asyncio.run(make_baselines_async())