Support functions as file plugins too

It is possible to write plugins which are only a function. At the moment they
are called on each line manually. This allows the function also to be called
on each file once. It works similar to creating the class and calling `run` on
it immediately. The plugin function needs to return a generator.

This is based on the original comment in the `FileChecker.run_ast_checks`
method, but slightly modified as the original comment would've called the
return of the function. But the function could return the reports directly.

2 files changed, 80 insertions, 6 deletions

diff --git a/src/flake8/checker.py b/src/flake8/checker.py
index ba412e3..99bacb5 100644
--- a/src/flake8/checker.py
+++ b/src/flake8/checker.py
@@ -466,11 +466,13 @@ class FileChecker(object):
 
         for plugin in self.checks.ast_plugins:
             checker = self.run_check(plugin, tree=ast)
-            # NOTE(sigmavirus24): If we want to allow for AST plugins that are
-            # not classes exclusively, we can do the following:
-            # retrieve_results = getattr(checker, 'run', lambda: checker)
-            # Otherwise, we just call run on the checker
-            for (line_number, offset, text, check) in checker.run():
+            # If the plugin uses a class, call the run method of it, otherwise
+            # the call should return something iterable itself
+            try:
+                runner = checker.run()
+            except AttributeError:
+                runner = checker
+            for (line_number, offset, text, check) in runner:
                 self.report(
                     error_code=None,
                     line_number=line_number,
diff --git a/tests/unit/test_plugin.py b/tests/unit/test_plugin.py
index f69bc05..73df659 100644
--- a/tests/unit/test_plugin.py
+++ b/tests/unit/test_plugin.py
@@ -4,7 +4,7 @@ import optparse
 import mock
 import pytest
 
-from flake8 import exceptions
+from flake8 import checker, exceptions
 from flake8.plugins import manager
 
 
@@ -151,3 +151,75 @@ def test_provide_options():
     plugin_obj.parse_options.assert_called_once_with(
         option_manager, option_values, None
     )
+
+
+EXPECTED_REPORT = (1, 1, 'T000 Expected Message')
+
+
+class PluginClass(object):
+    """Simple file plugin class yielding the expected report."""
+
+    name = 'test'
+    version = '1.0.0'
+
+    def __init__(self, tree):
+        """Dummy constructor to provide mandatory parameter."""
+        pass
+
+    def run(self):
+        """Run class yielding one element containing the expected report."""
+        yield EXPECTED_REPORT + (type(self), )
+
+
+def plugin_func(func):
+    """Decorator for file plugins which are implemented as functions."""
+    func.name = 'test'
+    func.version = '1.0.0'
+    return func
+
+
+@plugin_func
+def plugin_func_gen(tree):
+    """Simple file plugin function yielding the expected report."""
+    yield EXPECTED_REPORT + (type(plugin_func_gen), )
+
+
+@plugin_func
+def plugin_func_list(tree):
+    """Simple file plugin function returning a list of reports."""
+    return [EXPECTED_REPORT + (type(plugin_func_list), )]
+
+
+@pytest.mark.parametrize('plugin_target', [
+    PluginClass,
+    plugin_func_gen,
+    plugin_func_list,
+])
+def test_handle_file_plugins(plugin_target):
+    """Test the FileChecker class handling different file plugin types."""
+    # Mock an entry point returning the plugin target
+    entry_point = mock.Mock(spec=['require', 'resolve', 'load'])
+    entry_point.name = plugin_target.name
+    entry_point.resolve.return_value = plugin_target
+
+    # Load the checker plugins using the entry point mock
+    with mock.patch('pkg_resources.iter_entry_points',
+                    return_value=[entry_point]):
+        checks = manager.Checkers()
+
+    # Prevent it from reading lines from stdin or somewhere else
+    with mock.patch('flake8.processor.FileProcessor.read_lines',
+                    return_value=['Line 1']):
+        file_checker = checker.FileChecker('-', checks, mock.MagicMock())
+
+    # Do not actually build an AST
+    file_checker.processor.build_ast = lambda: True
+
+    # Forward reports to this mock
+    report = mock.Mock()
+    file_checker.report = report
+    file_checker.run_ast_checks()
+    report.assert_called_once_with(error_code=None,
+                                   line_number=EXPECTED_REPORT[0],
+                                   column=EXPECTED_REPORT[1],
+                                   text=EXPECTED_REPORT[2])