👌 IMPROVE: How mathjax is overriden

Only override if we have to, and override in a more precise manner.

myst_parser/__init__.py

@@ -18,17 +18,17 @@ def setup_sphinx(app):
     # we do this separately to setup,
     # so that it can be called by external packages like myst_nb
     from myst_parser.myst_refs import MystReferenceResolver
-    from myst_parser.myst_amsmath import MystAmsMathTransform
+    from myst_parser.mathjax import override_mathjax
     from myst_parser.main import MdParserConfig
 
     app.add_post_transform(MystReferenceResolver)
-    app.add_post_transform(MystAmsMathTransform)
 
     for name, default in MdParserConfig().as_dict().items():
         if not name == "renderer":
             app.add_config_value(f"myst_{name}", default, "env")
 
     app.connect("builder-inited", create_myst_config)
+    app.connect("builder-inited", override_mathjax)
 
 
 def create_myst_config(app):
@@ -50,27 +50,3 @@ def create_myst_config(app):
     except (TypeError, ValueError) as error:
         logger.error("myst configuration invalid: %s", error.args[0])
         app.env.myst_config = MdParserConfig()
-
-    # https://docs.mathjax.org/en/v2.7-latest/options/preprocessors/tex2jax.html#configure-tex2jax
-    if app.config.mathjax_config is None and app.env.myst_config.update_mathjax:
-        app.config.mathjax_config = {
-            "tex2jax": {
-                "inlineMath": [["\\(", "\\)"]],
-                "displayMath": [["\\[", "\\]"]],
-                "processRefs": False,
-                "processEnvironments": False,
-            }
-        }
-    elif app.env.myst_config.update_mathjax:
-        if "tex2jax" in app.config.mathjax_config:
-            logger.warning(
-                "`mathjax_config['tex2jax']` is set, but `myst_update_mathjax = True`, "
-                "and so this will be overridden. "
-                "Set `myst_update_mathjax = False` if you wish to use your own config"
-            )
-        app.config.mathjax_config["tex2jax"] = {
-            "inlineMath": [["\\(", "\\)"]],
-            "displayMath": [["\\[", "\\]"]],
-            "processRefs": False,
-            "processEnvironments": False,
-        }

myst_parser/mathjax.py

@@ -0,0 +1,83 @@
+"""Overrides to ``sphinx.ext.mathjax``
+
+This fixes two issues:
+
+1. Mathjax should not search for ``$`` delimiters, nor LaTeX amsmath environments,
+   since we already achieve this with the dollarmath and amsmath mrakdown-it-py plugins
+2. amsmath math blocks should be wrapped in mathjax delimiters (default ``\\[...\\]``),
+   and assigned an equation number
+
+"""
+from docutils import nodes
+
+from sphinx.application import Sphinx
+from sphinx.ext import mathjax
+from sphinx.locale import _
+from sphinx.util import logging
+from sphinx.util.math import get_node_equation_number
+from sphinx.writers.html import HTMLTranslator
+
+logger = logging.getLogger(__name__)
+
+
+def override_mathjax(app: Sphinx):
+    """Override aspects of the mathjax extension, but only if necessary."""
+
+    if (
+        app.config["myst_amsmath_enable"]
+        and "mathjax" in app.registry.html_block_math_renderers
+    ):
+        app.registry.html_block_math_renderers["mathjax"] = (
+            html_visit_displaymath,
+            None,
+        )
+    # https://docs.mathjax.org/en/v2.7-latest/options/preprocessors/tex2jax.html#configure-tex2jax
+    if app.config.mathjax_config is None and app.env.myst_config.update_mathjax:
+        app.config.mathjax_config = {
+            "tex2jax": {
+                "inlineMath": [["\\(", "\\)"]],
+                "displayMath": [["\\[", "\\]"]],
+                "processRefs": False,
+                "processEnvironments": False,
+            }
+        }
+    elif app.env.myst_config.update_mathjax:
+        if "tex2jax" in app.config.mathjax_config:
+            logger.warning(
+                "`mathjax_config['tex2jax']` is set, but `myst_update_mathjax = True`, "
+                "and so this will be overridden. "
+                "Set `myst_update_mathjax = False` if you wish to use your own config"
+            )
+        app.config.mathjax_config["tex2jax"] = {
+            "inlineMath": [["\\(", "\\)"]],
+            "displayMath": [["\\[", "\\]"]],
+            "processRefs": False,
+            "processEnvironments": False,
+        }
+
+
+def html_visit_displaymath(self: HTMLTranslator, node: nodes.math_block) -> None:
+    """Override for sphinx.ext.mathjax.html_visit_displaymath to handle amsmath.
+
+    By default displaymath, are normally wrapped in a prefix/suffix,
+    defined by mathjax_display, and labelled nodes are numbered.
+    However, this is not the case if the math_block is set as 'nowrap', as for amsmath.
+    Therefore, we need to override this behaviour.
+    """
+    if "amsmath" in node.get("classes", []):
+        self.body.append(
+            self.starttag(node, "div", CLASS="math notranslate nohighlight amsmath")
+        )
+        if node["number"]:
+            number = get_node_equation_number(self, node)
+            self.body.append('<span class="eqno">(%s)' % number)
+            self.add_permalink_ref(node, _("Permalink to this equation"))
+            self.body.append("</span>")
+        prefix, suffix = self.builder.config.mathjax_display
+        self.body.append(prefix)
+        self.body.append(self.encode(node.astext()))
+        self.body.append(suffix)
+        self.body.append("</div>\n")
+        raise nodes.SkipNode
+
+    return mathjax.html_visit_displaymath(self, node)

tests/test_sphinx/test_sphinx_builds/test_extended_syntaxes.html

@@ -31,12 +31,13 @@ <h1>
     <p>
      $ a=1 $
     </p>
-    <div class="amsmath" id="equation-mock-uuid">
+    <div class="amsmath math notranslate nohighlight" id="equation-mock-uuid">
      <span class="eqno">
       (2)
+      <a class="headerlink" href="#equation-mock-uuid" title="Permalink to this equation">
+       ¶
+      </a>
      </span>
-    </div>
-    <div class="amsmath math notranslate nohighlight" id="equation-mock-uuid">
      \[\begin{equation}
 b=2
 \end{equation}\]