From 6943ebad7c881f1ff8fc5d0a0a87b1174ea7360b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Mart=C3=ADnez=20Portela?=
 <sergio@codigoparallevar.com>
Date: Tue, 3 Aug 2021 22:37:03 +0200
Subject: [PATCH] Add support for org-roam-v2 style files.

---
 org_rw/org_rw.py         | 40 +++++++++++++++++++++++++++++++++++-----
 tests/07-org-roam-v2.org | 13 +++++++++++++
 tests/test_org.py        | 12 ++++++++++++
 3 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 tests/07-org-roam-v2.org

diff --git a/org_rw/org_rw.py b/org_rw/org_rw.py
index 47843bf..e26ee01 100644
--- a/org_rw/org_rw.py
+++ b/org_rw/org_rw.py
@@ -1261,7 +1261,9 @@ def parse_headline(hl, doc, parent) -> Headline:
 
 
 class OrgDoc:
-    def __init__(self, headlines, keywords, contents, list_items):
+    def __init__(
+        self, headlines, keywords, contents, list_items, structural, properties
+    ):
         self.todo_keywords = DEFAULT_TODO_KEYWORDS
         self.done_keywords = DEFAULT_DONE_KEYWORDS
 
@@ -1275,6 +1277,8 @@ class OrgDoc:
         self.keywords: List[Property] = keywords
         self.contents: List[RawLine] = contents
         self.list_items: List[ListItem] = list_items
+        self.structural: List = structural
+        self.properties: List = properties
         self._path = None
         self.headlines: List[Headline] = list(
             map(lambda hl: parse_headline(hl, self, self), headlines)
@@ -1292,13 +1296,20 @@ class OrgDoc:
         for content in self.contents:
             yield from get_links_from_content(content)
 
-    def get_property(self, name: str, default=None):
+    def get_keywords(self, name: str, default=None):
         for prop in self.keywords:
             if prop.key == name:
                 return prop.value
 
         return default
 
+    def get_property(self, name: str, default=None):
+        for prop in self.properties:
+            if prop.key == name:
+                return prop.value
+
+        return default
+
     def getProperties(self):
         return self.keywords
 
@@ -1453,6 +1464,12 @@ class OrgDoc:
 
     def dump(self):
         lines = []
+        for prop in self.properties:
+            lines.append(self.dump_property(prop))
+
+        for struct in self.structural:
+            lines.append(self.dump_structural(struct))
+
         for kw in self.keywords:
             lines.append(self.dump_kw(kw))
 
@@ -1476,9 +1493,18 @@ class OrgDocReader:
         self.contents: List[RawLine] = []
         self.delimiters: List[DelimiterLine] = []
         self.list_items: List[ListItem] = []
+        self.structural: List = []
+        self.properties: List = []
 
     def finalize(self):
-        return OrgDoc(self.headlines, self.keywords, self.contents, self.list_items)
+        return OrgDoc(
+            self.headlines,
+            self.keywords,
+            self.contents,
+            self.list_items,
+            self.structural,
+            self.properties,
+        )
 
     ## Construction
     def add_headline(self, linenum: int, match: re.Match) -> int:
@@ -1569,8 +1595,12 @@ class OrgDocReader:
             self.headline_hierarchy[-1]["delimiters"].append(line)
 
     def add_property_drawer_line(self, linenum: int, line: str, match: re.Match) -> int:
-        self.current_drawer = self.headline_hierarchy[-1]["properties"]
-        self.headline_hierarchy[-1]["structural"].append((linenum, line))
+        if len(self.headline_hierarchy) == 0:
+            self.current_drawer = self.properties
+            self.structural.append((linenum, line))
+        else:
+            self.current_drawer = self.headline_hierarchy[-1]["properties"]
+            self.headline_hierarchy[-1]["structural"].append((linenum, line))
 
     def add_results_drawer_line(self, linenum: int, line: str, match: re.Match) -> int:
         self.current_drawer = self.headline_hierarchy[-1]["results"]
diff --git a/tests/07-org-roam-v2.org b/tests/07-org-roam-v2.org
new file mode 100644
index 0000000..0707560
--- /dev/null
+++ b/tests/07-org-roam-v2.org
@@ -0,0 +1,13 @@
+:PROPERTIES:
+:ID:       515054a9-ced8-4119-a844-71726f80dedf
+:END:
+#+title: 07-org-roam-v2
+
+* Subnode
+:PROPERTIES:
+:CREATED:  [2021-08-03 Mar 22:12]
+:ID:       419f4651-21c8-4166-b8d5-692c34be9f93
+:END:
+
+** Sub sub node
+- This has no ID
diff --git a/tests/test_org.py b/tests/test_org.py
index 6bdda48..2a09c7c 100644
--- a/tests/test_org.py
+++ b/tests/test_org.py
@@ -488,3 +488,15 @@ class TestSerde(unittest.TestCase):
         self.assertEqual(lists2[1][1].content, " Second element")
         self.assertEqual(lists2[1][1].counter, "2")
         self.assertEqual(lists2[1][1].counter_sep, ")")
+
+    def test_org_roam_07(self):
+        with open(os.path.join(DIR, "07-org-roam-v2.org")) as f:
+            orig = f.read()
+            doc = loads(orig)
+
+        self.assertEqual(doc.get_property("ID"), "515054a9-ced8-4119-a844-71726f80dedf")
+
+        self.assertEqual(len(doc.getTopHeadlines()), 1)
+        hl = doc.getTopHeadlines()[0]
+        self.assertEqual(hl.get_property("ID"), "419f4651-21c8-4166-b8d5-692c34be9f93")
+        self.assertEqual(len(hl.children), 1)