feat(filetype): fall back to file extension when matching from hashbang (#22140)

If nothing matched in match_from_hashbang, also check the file extension table.
For a hashbang like '#!/bin/env foo', this will set the filetype to 'fooscript'
assuming the filetype for the 'foo' extension is 'fooscript' in the extension
table.

3 files changed, 24 insertions, 5 deletions

diff --git a/runtime/lua/vim/filetype.lua b/runtime/lua/vim/filetype.lua
index ec568b5b82..ca91f3f402 100644
--- a/runtime/lua/vim/filetype.lua
+++ b/runtime/lua/vim/filetype.lua
@@ -2604,7 +2604,9 @@ function M.match(args)
     -- If the function tries to use the filename that is nil then it will fail,
     -- but this enables checks which do not need a filename to still work.
     local ok
-    ok, ft = pcall(require('vim.filetype.detect').match_contents, contents, name)
+    ok, ft = pcall(require('vim.filetype.detect').match_contents, contents, name, function(ext)
+      return dispatch(extension[ext], name, bufnr)
+    end)
     if ok and ft then
       return ft
     end
diff --git a/runtime/lua/vim/filetype/detect.lua b/runtime/lua/vim/filetype/detect.lua
index edffdde9c7..b3d9fedeae 100644
--- a/runtime/lua/vim/filetype/detect.lua
+++ b/runtime/lua/vim/filetype/detect.lua
@@ -1420,7 +1420,7 @@ local patterns_hashbang = {
 
 ---@private
 -- File starts with "#!".
-local function match_from_hashbang(contents, path)
+local function match_from_hashbang(contents, path, dispatch_extension)
   local first_line = contents[1]
   -- Check for a line like "#!/usr/bin/env {options} bash".  Turn it into
   -- "#!/usr/bin/bash" to make matching easier.
@@ -1473,6 +1473,11 @@ local function match_from_hashbang(contents, path)
       return ft
     end
   end
+
+  -- If nothing matched, check the extension table. For a hashbang like
+  -- '#!/bin/env foo', this will set the filetype to 'fooscript' assuming
+  -- the filetype for the 'foo' extension is 'fooscript' in the extension table.
+  return dispatch_extension(name)
 end
 
 local patterns_text = {
@@ -1652,10 +1657,10 @@ local function match_from_text(contents, path)
   return cvs_diff(path, contents)
 end
 
-M.match_contents = function(contents, path)
+function M.match_contents(contents, path, dispatch_extension)
   local first_line = contents[1]
   if first_line:find('^#!') then
-    return match_from_hashbang(contents, path)
+    return match_from_hashbang(contents, path, dispatch_extension)
   else
     return match_from_text(contents, path)
   end
diff --git a/test/functional/lua/filetype_spec.lua b/test/functional/lua/filetype_spec.lua
index 034665f717..add69235b6 100644
--- a/test/functional/lua/filetype_spec.lua
+++ b/test/functional/lua/filetype_spec.lua
@@ -98,10 +98,22 @@ describe('vim.filetype', function()
   it('works with contents #22180', function()
     eq('sh', exec_lua [[
       -- Needs to be set so detect#sh doesn't fail
-      vim.g.ft_ignore_pat = "\\.\\(Z\\|gz\\|bz2\\|zip\\|tgz\\)$"
+      vim.g.ft_ignore_pat = '\\.\\(Z\\|gz\\|bz2\\|zip\\|tgz\\)$'
       return vim.filetype.match({ contents = { '#!/usr/bin/env bash' } })
     ]])
   end)
+
+  it('considers extension mappings when matching from hashbang', function()
+    eq('fooscript', exec_lua [[
+      vim.filetype.add({
+        extension = {
+          foo = 'fooscript',
+        }
+      })
+      return vim.filetype.match({ contents = { '#!/usr/bin/env foo' } })
+    ]])
+  end)
+
 end)
 
 describe('filetype.lua', function()