Lots of minor formatting fixes for embedded functions.

function.lua

@@ -182,16 +182,26 @@ local function add_indent_to_string (str, indent)
 end
 
 local function get_docs_from_function_body (func_body)
+	-- Finds the documentation lines of a function.
+	-- Also returns the remaining non-documentation lines.
+
 	assert(type(func_body) == 'string')
-	local doc_lines = {}
+
+	local doc_lines, not_doc_lines, still_docs  =  {}, {}, true
 	for line in func_body:gmatch('[^\n]+', true) do
-		if not line:match('^%s*$') then
+		if still_docs and not line:match('^%s*$') then
 			local line_text = line:match('^%s*%-%-%s*(.*)%s*$')
-			if not line_text then  break  end
+			if line_text then
 				doc_lines[#doc_lines+1] = line_text
+			else
+				still_docs = false
 			end
 		end
-	return table.concat(doc_lines, '\n')
+		if not still_docs then
+			not_doc_lines[#not_doc_lines+1] = line
+		end
+	end
+	return table.concat(doc_lines, '\n'), table.concat(not_doc_lines, '\n')
 end
 
 local function wrap_text (text, max_width)
@@ -248,11 +258,15 @@ return function (value, depth, l, format_value)
 
 	local function_params, function_body  =  nil, '...'
 
-    if info.defined_how == 'string' then
+	if info.defined_how == 'string' or not info.doc then
-        -- Function was defined as a string
+		local params_text,   body  =  get_function_paramlist_and_body(info)
-		local _, body  =  get_function_paramlist_and_body(info)
+		local documentation, body  =  get_docs_from_function_body(body)
+		body = body:match('^%s*(.-)%s*$')
 		if #body <= NR_CHARS_IN_LONG_FUNCTION_BODY and not body:find '\n' and not body:find(FUNCTION_KEYWORD_MATCH) then
-			function_body  =  body
+			if info.defined_how == 'string' then  function_body  =  body  end
+		end
+		if body then
+			info.doc  =  documentation ~= '' and documentation
 		end
 	end
 
@@ -299,14 +313,6 @@ return function (value, depth, l, format_value)
     end
 
 	-- Doc
-	if not info.doc then
-		local function_body  =  select(2, get_function_paramlist_and_body(info))
-		if function_body then
-			local documentation  =  get_docs_from_function_body(function_body)
-			info.doc  =  documentation ~= '' and documentation
-		end
-	end
-
     if info.doc then
 		l[#l+1] = '\n'
 		local indent  =  l.options.indent .. '-- '
@@ -363,7 +369,11 @@ return function (value, depth, l, format_value)
         l[#l+1] = '--]]'
     end
 
+	-- Ignore spacing and function body if it's a Λ string.
+	if function_body ~= '' then
 	    l[#l+1] = '\n'
 		l[#l+1] = indent
-    l[#l+1] = '...\nend'
+		l[#l+1] = function_body
+	end
+    l[#l+1] = '\nend'
 end

test/test_function.lua

@@ -188,14 +188,14 @@ format_test {
 
 format_test {
     name   = 'Can find docs from string-loaded function',
-    input  = loadstring 'return function ()\n--Hello\nend' (),
+    input  = loadstring 'return function ()\n-- Hello\nreturn 2\nend' (),
-    expect = 'function ()\n    -- Hello\n\n    ...\nend',
+    expect = 'function ()\n    -- Hello\n\n    return 2\nend',
 }
 
 format_test {
     name   = 'String-loaded functions without docs, won\'t display Source file comment',
-    input  = loadstring 'return function () end' (),
+    input  = loadstring 'return function (a) return a end' (),
-    expect = 'function ()\n    -- Loaded from string\n\n    ...\nend',
+    expect = 'function (a)\n    -- Loaded from string\n\n    return a\nend',
 }
 
 --------------------------------------------------------------------------------
@@ -309,6 +309,22 @@ format_test {
     expect = 'function (self) return self end',
 }
 
+format_test {
+    name   = 'Embed functions in single mode without documentation',
+    single = true,
+    adv_getlocal = true,
+    input  = loadstring('return function () -- Hi\nreturn 1234 end')(),
+    expect = 'function () return 1234 end',
+}
+
+format_test {
+    name   = 'Embed functions and ignore the whitespace after the documentation',
+    single = true,
+    adv_getlocal = true,
+    input  = loadstring('return function () -- Hi\n\n\t \t\t  return 1234 end')(),
+    expect = 'function () return 1234 end',
+}
+
 --[[ NOTE: This turned out to be hard. Requires traversal of AST.
 format_test {
     name = 'Embed correct function, when two are on the same line, but with different arguments 1',
@@ -339,10 +355,26 @@ format_test {
     name   = 'If embedding a function would include newline, ignore',
     single = true,
     adv_getlocal = true,
-    input  = loadstring('return function () -- Hi\n return 1234 end')(),
+    input  = loadstring('return function ()\n\ta = math.random()\nreturn a + 15 end')(),
     expect = 'function () ... end',
 }
 
+--------------------------------------------------------------------------------
+-- Elaborate embedded functions
+
+format_test {
+    name   = 'Embed functions when elaborate formatting',
+    adv_getlocal = true,
+    input  = loadstring('return function () return 1234 end')(),
+    expect = 'function ()\n    -- Loaded from string\n\n    return 1234\nend',
+}
+
+format_test {
+    name   = 'Embedding function without body, includes only the documentation',
+    input  = loadstring 'return function () end' (),
+    expect = 'function ()\n    -- Loaded from string\nend',
+}
+
 --------------------------------------------------------------------------------
 -- Indent functions nicely