Argument types for iframe

lib/streamlit/elements/iframe.py

@@ -12,20 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
-from typing import cast
+from typing import cast, Optional, TYPE_CHECKING
 
-import streamlit
 from streamlit.proto.IFrame_pb2 import IFrame as IFrameProto
 
+if TYPE_CHECKING:
+    from streamlit.delta_generator import DeltaGenerator
+
 
 class IframeMixin:
     def _iframe(
         self,
-        src,
-        width=None,
-        height=None,
-        scrolling=False,
+        src: str,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        scrolling: bool = False,
     ):
         """Load a remote URL in an iframe.
 
@@ -55,10 +56,10 @@ def _iframe(
 
     def _html(
         self,
-        html,
-        width=None,
-        height=None,
-        scrolling=False,
+        html: str,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        scrolling: bool = False,
     ):
         """Display an HTML string in an iframe.
 
@@ -87,13 +88,13 @@ def _html(
         return self.dg._enqueue("iframe", iframe_proto)
 
     @property
-    def dg(self) -> "streamlit.delta_generator.DeltaGenerator":
+    def dg(self) -> "DeltaGenerator":
         """Get our DeltaGenerator."""
-        return cast("streamlit.delta_generator.DeltaGenerator", self)
+        return cast("DeltaGenerator", self)
 
 
 def marshall(
-    proto,
+    proto: IFrameProto,
     src: Optional[str] = None,
     srcdoc: Optional[str] = None,
     width: Optional[int] = None,