X-Git-Url: https://git.mdrn.pl/librarian.git/blobdiff_plain/cec81dc04a23808077505265a3f8bb0e9b6053cb..63cfa40372a9d33c84c0e4b61ed505173c780eae:/librarian/fb2.py

diff --git a/librarian/fb2.py b/librarian/fb2.py
index b0ad410..1e110f5 100644
--- a/librarian/fb2.py
+++ b/librarian/fb2.py
@@ -7,10 +7,33 @@ import os.path
 from copy import deepcopy
 from lxml import etree
 
-from librarian import functions, OutputFile
+from librarian import functions, IOFile
+from .epub import replace_by_verse
 
 
 functions.reg_substitute_entities()
+functions.reg_person_name()
+
+
+def sectionify(tree):
+    """Finds section headers and adds a tree of _section tags."""
+    sections = ['naglowek_czesc',
+            'naglowek_akt', 'naglowek_rozdzial', 'naglowek_scena',
+            'naglowek_podrozdzial']
+    section_level = dict((v,k) for (k,v) in enumerate(sections))
+
+    # We can assume there are just subelements an no text at section level.
+    for level, section_name in reversed(list(enumerate(sections))):
+        for header in tree.findall('//' + section_name):
+            section = header.makeelement("_section")
+            header.addprevious(section)
+            section.append(header)
+            sibling = section.getnext()
+            while (sibling is not None and
+                    section_level.get(sibling.tag, 1000) > level):
+                section.append(sibling)
+                sibling = section.getnext()
+
 
 def transform(wldoc, verbose=False,
               cover=None, flags=None):
@@ -30,6 +53,11 @@ def transform(wldoc, verbose=False,
     style_filename = os.path.join(os.path.dirname(__file__), 'fb2/fb2.xslt')
     style = etree.parse(style_filename)
 
+    replace_by_verse(document.edoc)
+    sectionify(document.edoc)
+
     result = document.transform(style)
 
-    return OutputFile.from_string(unicode(result).encode('utf-8'))
+    return IOFile.from_string(unicode(result).encode('utf-8'))
+
+# vim:et