From 66ccb662117aa034365dcef470f74738d52738a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Mart=C3=ADnez=20Portela?=
 <sergio@codigoparallevar.com>
Date: Mon, 30 Sep 2024 23:39:43 +0200
Subject: [PATCH] Explicitly extract code block language.

---
 org_rw/org_rw.py  | 10 +++++++++-
 tests/test_org.py |  5 ++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/org_rw/org_rw.py b/org_rw/org_rw.py
index c5ae4f7..8cb3855 100644
--- a/org_rw/org_rw.py
+++ b/org_rw/org_rw.py
@@ -113,7 +113,7 @@ BEGIN_BLOCK_RE = re.compile(r"^\s*#\+BEGIN_(?P<subtype>[^ ]+)(?P<arguments>.*)$"
 END_BLOCK_RE = re.compile(r"^\s*#\+END_(?P<subtype>[^ ]+)\s*$", re.I)
 RESULTS_DRAWER_RE = re.compile(r"^\s*:results:\s*$", re.I)
 CodeSnippet = collections.namedtuple(
-    "CodeSnippet", ("name", "content", "result", "arguments")
+    "CodeSnippet", ("name", "content", "result", "language", "arguments")
 )
 
 # Groupings
@@ -916,12 +916,19 @@ class Headline:
                     # the content parsing must be re-thinked
                     contents = contents[:-1]
 
+                language = None
+                if arguments is not None:
+                    arguments = arguments.strip()
+                    if " " in arguments:
+                        language = arguments[: arguments.index(" ")]
+                        arguments = arguments[arguments.index(" ") + 1 :]
                 sections.append(
                     {
                         "line_first": start + 1,
                         "line_last": end - 1,
                         "content": contents,
                         "arguments": arguments,
+                        "language": language,
                         "name": name,
                     }
                 )
@@ -983,6 +990,7 @@ class Headline:
                     content=content,
                     result=code_result,
                     arguments=arguments,
+                    language=language,
                     name=name,
                 )
             )
diff --git a/tests/test_org.py b/tests/test_org.py
index 5a0bc53..ad35b89 100644
--- a/tests/test_org.py
+++ b/tests/test_org.py
@@ -481,6 +481,7 @@ class TestSerde(unittest.TestCase):
         snippets = list(doc.get_code_snippets())
         self.assertEqual(len(snippets), 3)
         self.assertEqual(snippets[0].name, "first-code-name")
+        self.assertEqual(snippets[0].language, "shell")
         self.assertEqual(
             snippets[0].content,
             'echo "This is a test"\n'
@@ -488,7 +489,7 @@ class TestSerde(unittest.TestCase):
             + "exit 0 # Exit successfully",
         )
         self.assertEqual(
-            snippets[0].arguments.split(), ["shell", ":results", "verbatim"]
+            snippets[0].arguments.split(), [":results", "verbatim"]
         )
         self.assertEqual(
             snippets[0].result,
@@ -496,6 +497,7 @@ class TestSerde(unittest.TestCase):
         )
 
         self.assertEqual(snippets[1].name, None)
+        self.assertEqual(snippets[1].language, "shell")
         self.assertEqual(
             snippets[1].content,
             'echo "This is another test"\n'
@@ -507,6 +509,7 @@ class TestSerde(unittest.TestCase):
         )
 
         self.assertEqual(snippets[2].name, None)
+        self.assertEqual(snippets[2].language, "c")
         self.assertEqual(
             snippets[2].content,
             "/* This code has to be escaped to\n"