Chunks are now formatted correctly.

function.lua

@@ -195,17 +195,30 @@ local function get_function_body_info (info)
 
 	if not str then  return info  end
 
+	-- Now find some info about the function.
+	-- NOTE: function_params is currently not used for anything.
+	local function_name, function_params, function_body
+
+	if info.linedefined == 0 then
+		-- A function is a "chunk" when linedefined is equal to 0.
+		-- This is a "toplevel" function. One without a function
+		-- definition. The entire string is it's function body.
+		function_name	= ''
+		function_body	= str
+	else
+		-- This is not a chunk. Look for function definition.
+
 		-- Calculate indices of the lines the function should be defined at.
 		local start_line_index	=  get_line_index(str, info.linedefined)
 		local end_line_index	=  get_line_index(str, info.lastlinedefined + 1)
 
-	-- Now find some info about the function.
-	-- NOTE: function_params is currently not used for anything.
-	local function_name, function_params, function_body  =  str:sub(start_line_index, end_line_index):match(FUNCTION_DEFINITION_MATCH)
+		function_name, function_params, function_body  =  str:sub(start_line_index, end_line_index):match(FUNCTION_DEFINITION_MATCH)
 
+		-- Throw an error if we can't find anything.
 		if type(function_body) ~= 'string' then
 			error(('[pretty.function/internal]: Could not find the function defined on lines %i-%i (indices %i-%i) for string:\n\n%s\n'):format(info.linedefined, info.lastlinedefined, start_line_index, end_line_index, str))
 		end
+	end
 
 	local function_body	=  function_body:match '^%s*(.-)%s*$'
 	local pivot_index	=  get_docs_split_index (function_body)

test/test_function.lua

@@ -269,6 +269,20 @@ format_test {
     expect = 'function (a, b) return a + b end',
 }
 
+format_test {
+    name   = 'Embedding loaded chunk',
+    single = true,
+    input  = loadstring 'return 42',
+    expect = 'function (...) return 42 end',
+}
+
+format_test {
+    name   = 'Embedding loaded file',
+    single = true,
+    input  = loadfile 'init.lua',
+    expect = 'function (...) ... end',
+}
+
 format_test {
     name = 'Embedding nested function, when on same line is too hard 1',
     single = true,

test/test_pretty.lua

@@ -341,12 +341,6 @@ SUITE:addTest('UseCase: Can load function from file that is shortly deleted', fu
     assert(true)
 end)
 
-SUITE:addTest('UseCase: Can use pretty on loadstring block', function ()
-    -- TODO: Move to more appropriate test module.
-    format(loadstring 'hi = 0')
-    assert(true)
-end)
-
 --------------------------------------------------------------------------------
 
 return SUITE