Expand Markdown detection to all Python language names (#268)

Author: adamchainz

CHANGELOG.rst

@@ -6,6 +6,8 @@ Changelog
 
   Thanks to initial work from Matthew Anderson in `PR #246 <https://github.com/adamchainz/blacken-docs/pull/246>`__.
 
+* Expand Markdown detection to all Python language names from Pygments: ``py``, ``sage``, ``python3``, ``py3``, and ``numpy``.
+
 * Remove ``language_version`` from ``.pre-commit-hooks.yaml``.
   This change allows ``default_language_version`` in ``.pre-commit-config.yaml` to take precedence.
 

src/blacken_docs/__init__.py

@@ -13,8 +13,10 @@
 from black.mode import TargetVersion
 
 
+PYGMENTS_PY_LANGS = frozenset(("python", "py", "sage", "python3", "py3", "numpy"))
+PYGMENTS_PY_LANGS_RE_FRAGMENT = f"({'|'.join(PYGMENTS_PY_LANGS)})"
 MD_RE = re.compile(
-    r"(?P<before>^(?P<indent> *)```\s*python( .*?)?\n)"
+    r"(?P<before>^(?P<indent> *)```\s*" + PYGMENTS_PY_LANGS_RE_FRAGMENT + r"( .*?)?\n)"
     r"(?P<code>.*?)"
     r"(?P<after>^(?P=indent)```\s*$)",
     re.DOTALL | re.MULTILINE,
@@ -25,7 +27,6 @@
     r"(?P<after>^(?P=indent)```.*$)",
     re.DOTALL | re.MULTILINE,
 )
-RST_PY_LANGS = frozenset(("python", "py", "sage", "python3", "py3", "numpy"))
 BLOCK_TYPES = "(code|code-block|sourcecode|ipython)"
 DOCTEST_TYPES = "(testsetup|testcleanup|testcode)"
 RST_RE = re.compile(
@@ -117,7 +118,7 @@ def _md_match(match: Match[str]) -> str:
 
     def _rst_match(match: Match[str]) -> str:
         lang = match["lang"]
-        if lang is not None and lang not in RST_PY_LANGS:
+        if lang is not None and lang not in PYGMENTS_PY_LANGS:
             return match[0]
         min_indent = min(INDENT_RE.findall(match["code"]))
         trailing_ws_match = TRAILING_NL_RE.search(match["code"])

tests/test_blacken_docs.py

@@ -29,6 +29,12 @@ def test_format_src_markdown_leading_whitespace():
     assert after == ("```   python\n" "f(1, 2, 3)\n" "```\n")
 
 
+def test_format_src_markdown_short_name():
+    before = "```   py\n" "f(1,2,3)\n" "```\n"
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == ("```   py\n" "f(1, 2, 3)\n" "```\n")
+
+
 def test_format_src_markdown_options():
     before = "```python title='example.py'\n" + "f(1,2,3)\n" + "```\n"
     after, _ = blacken_docs.format_str(before, BLACK_MODE)