Support python related languages in to_notion_language

pyproject.toml

@@ -49,6 +49,7 @@ dev = [
     "flake8",
     "mypy",
     "types-docutils",
+    "types-Pygments",
 ]
 
 [tool.taskipy.tasks]

src/sphinx_notion/nodes/literal_block.py

@@ -0,0 +1,23 @@
+from pygments.lexers import get_lexer_by_name
+
+PygmentsLanguage = str
+
+
+def get_standard_pygments_language(language: str) -> PygmentsLanguage:
+    # "default" language means "not specified"
+    if language == "default":
+        return "default"
+    lexer = get_lexer_by_name(language)
+    # Lexer has aliases but mypy raises this error:
+    # >error: "Lexer" has no attribute "aliases"  [attr-defined]
+    # https://github.com/pygments/pygments/blob/2.19.2/pygments/lexer.py#L111-L113
+    return lexer.aliases[0]  # type: ignore[attr-defined]
+
+
+def to_notion_language(pygments_language: PygmentsLanguage) -> str:
+    if pygments_language in {"default", "pytb", "text", "output"}:
+        return "plain text"
+    if pygments_language in {"python", "pycon"}:
+        return "python"
+    # TODO: Support for other languages
+    return pygments_language

src/sphinx_notion/writers.py

@@ -7,14 +7,10 @@
 from sphinx.builders.text import TextBuilder
 from sphinx.writers.text import TextTranslator
 
-
-def to_notion_language(pygments_language: str) -> str:
-    if pygments_language == "default":
-        # default means "not specified"
-        return "plain text"
-    if pygments_language == "text":
-        return "plain text"
-    return pygments_language
+from sphinx_notion.nodes.literal_block import (
+    get_standard_pygments_language,
+    to_notion_language,
+)
 
 
 class NotionTranslator(TextTranslator):
@@ -125,6 +121,9 @@ def visit_bullet_list(self, node: nodes.Element) -> None:
     def visit_literal_block(self, node: nodes.Element) -> None:
         super().visit_literal_block(node)
 
+        pygments_language = get_standard_pygments_language(
+            node.attributes["language"]
+        )
         self._json.append(
             {
                 "object": "block",
@@ -133,9 +132,7 @@ def visit_literal_block(self, node: nodes.Element) -> None:
                     "rich_text": [
                         {"type": "text", "text": {"content": node.astext()}}
                     ],
-                    "language": to_notion_language(
-                        node.attributes["language"]
-                    ),
+                    "language": to_notion_language(pygments_language),
                 },
             }
         )

tests/nodes/test_literal_block.py

@@ -0,0 +1,39 @@
+import pytest
+
+from sphinx_notion.nodes.literal_block import (
+    get_standard_pygments_language,
+    to_notion_language,
+)
+
+
+@pytest.mark.parametrize(
+    "language, expected",
+    [
+        ("python", "python"),
+        ("python3", "python"),
+        ("pycon", "pycon"),
+        ("python-console", "pycon"),
+        ("pytb", "pytb"),
+        ("py3tb", "pytb"),
+        ("text", "text"),
+        ("output", "output"),
+        ("default", "default"),
+    ],
+)
+def test_get_standard_pygments_language(language, expected):
+    assert get_standard_pygments_language(language) == expected
+
+
+@pytest.mark.parametrize(
+    "pygments_language, expected",
+    [
+        ("python", "python"),
+        ("default", "plain text"),
+        ("pycon", "python"),
+        ("pytb", "plain text"),
+        ("text", "plain text"),
+        ("output", "plain text"),
+    ],
+)
+def test_to_notion_language(pygments_language, expected):
+    assert to_notion_language(pygments_language) == expected

tests/roots/test-code/expected.json

@@ -22,7 +22,7 @@
                 {
                     "type": "text",
                     "text": {
-                        "content": "print(3 * 4)"
+                        "content": "print(5 - 6)"
                     }
                 }
             ],
@@ -37,11 +37,11 @@
                 {
                     "type": "text",
                     "text": {
-                        "content": "print(5 - 6)"
+                        "content": ">>> print(\"Hello, World!\")\nHello, World!"
                     }
                 }
             ],
-            "language": "plain text"
+            "language": "python"
         }
     }
 ]

tests/roots/test-code/source/index.rst

@@ -3,10 +3,11 @@
     print(1 + 2)
     print("Hello")
 
-.. code-block::
-
-    print(3 * 4)
-
 .. code-block:: text
 
     print(5 - 6)
+
+.. code-block:: python-console
+
+    >>> print("Hello, World!")
+    Hello, World!