stubgen removes Protocol from base classes list

[citruz]

Bug Report

Consider the following input:

from typing import Protocol

class Interface(Protocol):
    def dummy(self, arg: str) -> str:
        ...

stubgen produces:

class Interface:
    def dummy(self, arg: str) -> str: ...

Note that the Protocol class is removed from the list of base classes.

Obviously this causes problems when performing type checking against users of that generated typestub:

from interface import Interface

class InterfaceImpl:
    def dummy(self, arg: str) -> str:
        return arg


interf: Interface = InterfaceImpl()

mypy:

provider/__init__.py:8: error: Incompatible types in assignment (expression has type "InterfaceImpl", variable has type "Interface")

Expected Behavior

Protocol base class is not removed.

Actual Behavior

Protocol base class is removed.

Your Environment

• Mypy version used: mypy 0.931
• Python version used: Python 3.9.7
• Operating system and version: macOS

[citruz]

This works as a fix:

diff --git a/mypy/stubgen.py b/mypy/stubgen.py
index b0a2c4e64..4e1a0fe08 100755
--- a/mypy/stubgen.py
+++ b/mypy/stubgen.py
@@ -849,6 +849,9 @@ class StubGenerator(mypy.traverser.TraverserVisitor):
             base_types.append('metaclass=abc.ABCMeta')
             self.import_tracker.add_import('abc')
             self.import_tracker.require_name('abc')
+        elif self.analyzed and o.info.is_protocol:
+            base_types.append('Protocol')
+            self.add_typing_import('Protocol')
         if base_types:
             self.add('(%s)' % ', '.join(base_types))
         self.add(':\n')

If you think this is the way to go, I can submit it as a PR.