graphframes · rjurney · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025 · Jan 16, 2025
diff --git a/python/docs/epytext.py b/python/docs/epytext.py
@@ -1,4 +1,5 @@
 import re
+from sphinx.application import Sphinx
 
 RULES = (
     (r"<(!BLANKLINE)[\w.]+>", r""),
@@ -9,7 +10,7 @@
     ('pyspark.rdd.RDD', 'RDD'),
 )
 
-def _convert_epytext(line):
+def _convert_epytext(line: str) -> str:
     """
     >>> _convert_epytext("L{A}")
     :class:`A`
@@ -19,9 +20,9 @@ def _convert_epytext(line):
         line = re.sub(p, sub, line)
     return line
 
-def _process_docstring(app, what, name, obj, options, lines):
+def _process_docstring(app: "Sphinx", what: str, name: str, obj: object, options: dict, lines: list[str]) -> None:
     for i in range(len(lines)):
         lines[i] = _convert_epytext(lines[i])
 
-def setup(app):
+def setup(app: "Sphinx") -> None:
     app.connect("autodoc-process-docstring", _process_docstring)
diff --git a/python/docs/underscores.py b/python/docs/underscores.py
@@ -29,27 +29,33 @@
 """
 import os
 import shutil
+from collections.abc import Callable
+from typing import Any
+from sphinx.application import Sphinx
 
-
-def setup(app):
+def setup(app: Sphinx) -> None:
     """
     Add a html-page-context  and a build-finished event handlers
     """
     app.connect('html-page-context', change_pathto)
     app.connect('build-finished', move_private_folders)
 
-def change_pathto(app, pagename, templatename, context, doctree):
+def change_pathto(app: Sphinx,
+                 pagename: str,
+                 templatename: str,
+                 context: dict[str, Any],
+                 doctree: Any | None) -> None:
     """
     Replace pathto helper to change paths to folders with a leading underscore.
     """
-    pathto = context.get('pathto')
-    def gh_pathto(otheruri, *args, **kw):
+    pathto: Callable = context.get('pathto')
+    def gh_pathto(otheruri: str, *args: Any, **kw: Any) -> Any:
         if otheruri.startswith('_'):
             otheruri = otheruri[1:]
         return pathto(otheruri, *args, **kw)
     context['pathto'] = gh_pathto
 
-def move_private_folders(app, e):
+def move_private_folders(app: Sphinx, e: Exception | None) -> None:
     """
     remove leading underscore from folders in in the output folder.
 

diff --git a/python/graphframes/examples/belief_propagation.py b/python/graphframes/examples/belief_propagation.py
@@ -16,6 +16,7 @@
 #
 
 import math
+from typing import Union
 
 # Import subpackage examples here explicitly so that
 # this module can be run directly with spark-submit.
@@ -27,7 +28,7 @@
 __all__ = ['BeliefPropagation']
 
 
-class BeliefPropagation(object):
+class BeliefPropagation:
     """Example code for Belief Propagation (BP)
 
     This provides a template for building customized BP algorithms for different types of graphical
@@ -63,7 +64,7 @@ class BeliefPropagation(object):
      """
 
     @classmethod
-    def runBPwithGraphFrames(cls, g, numIter):
+    def runBPwithGraphFrames(cls, g: GraphFrame, numIter: int) -> GraphFrame:
         """Run Belief Propagation using GraphFrame.
 
         This implementation of BP shows how to use GraphFrame's aggregateMessages method.
@@ -117,7 +118,7 @@ def runBPwithGraphFrames(cls, g, numIter):
         return GraphFrame(gx.vertices.drop('color'), gx.edges)
 
     @staticmethod
-    def _colorGraph(g):
+    def _colorGraph(g: GraphFrame) -> GraphFrame:
         """Given a GraphFrame, choose colors for each vertex.
 
         No neighboring vertices will share the same color. The number of colors is minimized.
@@ -135,7 +136,7 @@ def _colorGraph(g):
         return GraphFrame(v, g.edges)
 
     @staticmethod
-    def _sigmoid(x):
+    def _sigmoid(x: Union[int, float, None]) -> Union[float, None]:
         """Numerically stable sigmoid function 1 / (1 + exp(-x))"""
         if not x:
             return None
@@ -147,7 +148,7 @@ def _sigmoid(x):
             return z / (1 + z)
 
 
-def main():
+def main() -> None:
     """Run the belief propagation algorithm for an example problem."""
     # setup spark session
     spark = SparkSession.builder.appName("BeliefPropagation example").getOrCreate()

diff --git a/python/graphframes/examples/graphs.py b/python/graphframes/examples/graphs.py
@@ -17,24 +17,24 @@
 
 import itertools
 
-from pyspark.sql import functions as sqlfunctions
+from pyspark.sql import functions as sqlfunctions, SparkSession
 
 from graphframes import GraphFrame
 
 __all__ = ['Graphs']
 
 
-class Graphs(object):
+class Graphs:
     """Example GraphFrames for testing the API
 
     :param spark: SparkSession
     """
 
-    def __init__(self, spark):
+    def __init__(self, spark: SparkSession) -> None:
         self._spark = spark
         self._sc = spark._sc
 
-    def friends(self):
+    def friends(self) -> GraphFrame:
         """A GraphFrame of friends in a (fake) social network."""
         # Vertex DataFrame
         v = self._spark.createDataFrame([
@@ -58,7 +58,7 @@ def friends(self):
         # Create a GraphFrame
         return GraphFrame(v, e)
 
-    def gridIsingModel(self, n, vStd=1.0, eStd=1.0):
+    def gridIsingModel(self, n: int, vStd: float = 1.0, eStd: float = 1.0) -> GraphFrame:
         """Grid Ising model with random parameters.
 
         Ising models are probabilistic graphical models over binary variables x\ :sub:`i`.