Improve section detection, fixes #68

1 files changed, 9 insertions, 5 deletions

diff --git a/lib/parse/lex.js b/lib/parse/lex.js
index def74b4..cec6047 100644
--- a/lib/parse/lex.js
+++ b/lib/parse/lex.js
@@ -19,13 +19,17 @@ function splitSections(nodes) {
 
 // What is the type of this section
 function sectionType(nodes, idx) {
-    var codeNodes = _.filter(nodes, {
-        type: 'code'
-    }).length;
+    var codeType = { type: 'code' };
+
+    // Number of code nodes in section
+    var len = _.filter(nodes, codeType).length;
 
     if(
-        (codeNodes === 3 || codeNodes === 4) &&
-        (idx % 2) === 1)
+        // Got 3 or 4 code blocks
+        (len === 3 || len === 4) &&
+        // Ensure all nodes are at the end
+        _.all(_.last(nodes, len), codeType)
+    )
     {
         return 'exercise';
     }