diff --git a/CHANGELOG.md b/CHANGELOG.md index 2786b3e1ed..71b8eea74c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,12 @@ All notable changes to the PyGraphistry are documented in this file. The PyGraph The changelog format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/). This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html) and all PyGraphistry-specific breaking changes are explictly noted here. +## Dev + +### Breaking +* Plottable is now a Protocol +* py.typed added, type checking active on PyGraphistry! + ## [0.38.0 - 2025-06-17] ### Changed diff --git a/docs/source/api/ai.rst b/docs/source/api/ai.rst index a8ad9ea41b..7460dcd0b5 100644 --- a/docs/source/api/ai.rst +++ b/docs/source/api/ai.rst @@ -58,5 +58,4 @@ HeterographEmbedModuleMixin .. automodule:: graphistry.embed_utils :members: :undoc-members: - :inherited-members: :show-inheritance: diff --git a/docs/source/api/compute.rst b/docs/source/api/compute.rst index 5fbf09cc54..fb361e6ca2 100644 --- a/docs/source/api/compute.rst +++ b/docs/source/api/compute.rst @@ -6,7 +6,6 @@ ComputeMixin module .. automodule:: graphistry.compute.ComputeMixin :members: :undoc-members: - :inherited-members: :show-inheritance: Collapse diff --git a/docs/source/api/layout/index.rst b/docs/source/api/layout/index.rst index defa7a14eb..e9ad602be9 100644 --- a/docs/source/api/layout/index.rst +++ b/docs/source/api/layout/index.rst @@ -38,5 +38,4 @@ LayoutsMixin .. automodule:: graphistry.layouts.LayoutsMixin :members: :undoc-members: - :inherited-members: :show-inheritance: \ No newline at end of file diff --git a/docs/source/api/plotter.rst b/docs/source/api/plotter.rst index 361135c339..140d888cb1 100644 --- a/docs/source/api/plotter.rst +++ b/docs/source/api/plotter.rst @@ -18,18 +18,20 @@ Plotter Class :inherited-members: :show-inheritance: -Plottable Interface + +PlotterBase Class ---------------------- -.. automodule:: graphistry.Plottable +.. automodule:: graphistry.PlotterBase :members: :undoc-members: :inherited-members: :show-inheritance: -PlotterBase Class + +Plottable Interface ---------------------- -.. automodule:: graphistry.PlotterBase +.. automodule:: graphistry.Plottable :members: :undoc-members: :inherited-members: - :show-inheritance: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/api/plugins/compute/cugraph.rst b/docs/source/api/plugins/compute/cugraph.rst index 775ced993e..b49bad469f 100644 --- a/docs/source/api/plugins/compute/cugraph.rst +++ b/docs/source/api/plugins/compute/cugraph.rst @@ -9,7 +9,6 @@ cuGraph is a GPU-accelerated graph library that leverages the Nvidia RAPIDS ecos .. automodule:: graphistry.plugins.cugraph :members: :undoc-members: - :inherited-members: :show-inheritance: .. rubric:: Constants diff --git a/docs/source/api/plugins/compute/graphviz.rst b/docs/source/api/plugins/compute/graphviz.rst index 81c760aa4f..099abdab06 100644 --- a/docs/source/api/plugins/compute/graphviz.rst +++ b/docs/source/api/plugins/compute/graphviz.rst @@ -9,7 +9,6 @@ graphviz is a popular graph visualization library that PyGraphistry can interfac .. automodule:: graphistry.plugins.graphviz :members: :undoc-members: - :inherited-members: :show-inheritance: .. rubric:: Constants diff --git a/docs/source/api/plugins/compute/igraph.rst b/docs/source/api/plugins/compute/igraph.rst index d887e0e096..0290b43f61 100644 --- a/docs/source/api/plugins/compute/igraph.rst +++ b/docs/source/api/plugins/compute/igraph.rst @@ -9,7 +9,6 @@ igraph is a popular graph library that PyGraphistry can interface with. This all .. automodule:: graphistry.plugins.igraph :members: :undoc-members: - :inherited-members: :show-inheritance: .. rubric:: Constants diff --git a/docs/source/api/plugins/db/cosmos.rst b/docs/source/api/plugins/db/cosmos.rst index cacb7930aa..31075541a3 100644 --- a/docs/source/api/plugins/db/cosmos.rst +++ b/docs/source/api/plugins/db/cosmos.rst @@ -10,5 +10,4 @@ Azure Cosmos DB supports Gremlin graph queries .. autoclass:: graphistry.gremlin.CosmosMixin :members: :undoc-members: - :inherited-members: :show-inheritance: \ No newline at end of file diff --git a/docs/source/api/plugins/db/gremlin.rst b/docs/source/api/plugins/db/gremlin.rst index 25f8478bd9..ac54c20b27 100644 --- a/docs/source/api/plugins/db/gremlin.rst +++ b/docs/source/api/plugins/db/gremlin.rst @@ -12,5 +12,4 @@ As an open source technology, multiple databases support it. .. autoclass:: graphistry.gremlin.GremlinMixin :members: :undoc-members: - :inherited-members: :show-inheritance: \ No newline at end of file diff --git a/docs/source/api/plugins/db/neptune.rst b/docs/source/api/plugins/db/neptune.rst index c0687e3592..c50a9b3b1f 100644 --- a/docs/source/api/plugins/db/neptune.rst +++ b/docs/source/api/plugins/db/neptune.rst @@ -10,5 +10,4 @@ Amazon Neptune is a managed graph database by Amazon. It supports OpenCypher, RD .. autoclass:: graphistry.gremlin.NeptuneMixin :members: :undoc-members: - :inherited-members: :show-inheritance: diff --git a/graphistry/Engine.py b/graphistry/Engine.py index d2cd2943ea..aedf0fb10b 100644 --- a/graphistry/Engine.py +++ b/graphistry/Engine.py @@ -39,7 +39,6 @@ def resolve_engine( return Engine(engine.value) if g_or_df is not None: - # work around circular dependency from graphistry.Plottable import Plottable if isinstance(g_or_df, Plottable): if g_or_df._nodes is not None and g_or_df._edges is not None: @@ -48,7 +47,6 @@ def resolve_engine( warnings.warn(f'Edges and nodes must be same type for auto engine selection, got: {type(g_or_df._edges)} and {type(g_or_df._nodes)}') g_or_df = g_or_df._edges if g_or_df._edges is not None else g_or_df._nodes - if g_or_df is not None: if isinstance(g_or_df, pd.DataFrame): return Engine.PANDAS diff --git a/graphistry/Plottable.py b/graphistry/Plottable.py index 94c5c7ae9a..052c3d339b 100644 --- a/graphistry/Plottable.py +++ b/graphistry/Plottable.py @@ -1,18 +1,20 @@ -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union -from typing_extensions import Literal +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set, Tuple, Union, Protocol, overload +from typing_extensions import Literal, runtime_checkable import pandas as pd from graphistry.models.ModelDict import ModelDict from graphistry.models.compute.chain_remote import FormatType, OutputTypeAll, OutputTypeDf, OutputTypeGraph from graphistry.models.compute.dbscan import DBSCANEngine from graphistry.models.compute.umap import UMAPEngineConcrete +from graphistry.models.compute.features import GraphEntityKind from graphistry.plugins_types.cugraph_types import CuGraphKind +from graphistry.plugins_types.graphviz_types import EdgeAttr, Format, GraphAttr, NodeAttr, Prog from graphistry.plugins_types.kusto_types import KustoConfig from graphistry.plugins_types.spanner_types import SpannerConfig +from graphistry.privacy import Mode as PrivacyMode from graphistry.Engine import EngineAbstract from graphistry.utils.json import JSONVal - if TYPE_CHECKING: try: from umap import UMAP @@ -22,9 +24,20 @@ from sklearn.pipeline import Pipeline except: Pipeline = Any + try: + from cugraph import Graph + from cugraph import MultiGraph + from cugraph import BiPartiteGraph + except: + Graph = Any + MultiGraph = Any + BiPartiteGraph = Any else: UMAP = Any Pipeline = Any + Graph = Any + MultiGraph = Any + BiPartiteGraph = Any RenderModesConcrete = Literal["g", "url", "ipython", "databricks", "browser"] @@ -32,7 +45,8 @@ RenderModes = Union[Literal["auto"], RenderModesConcrete] RENDER_MODE_VALUES: Set[RenderModes] = set(["auto", "g", "url", "ipython", "databricks", "browser"]) -class Plottable(object): +@runtime_checkable +class Plottable(Protocol): _edges : Any _nodes : Any @@ -67,7 +81,7 @@ class Plottable(object): _complex_encodings : dict _bolt_driver : Any _tigergraph : Any - + _spanner_config: Optional[SpannerConfig] _kusto_config: Optional[KustoConfig] @@ -81,12 +95,14 @@ class Plottable(object): _node_features : Optional[pd.DataFrame] _node_features_raw: Optional[pd.DataFrame] _node_target : Optional[pd.DataFrame] + _node_target_raw : Optional[pd.DataFrame] _edge_embedding : Optional[pd.DataFrame] _edge_encoder : Optional[Any] _edge_features : Optional[pd.DataFrame] _edge_features_raw: Optional[pd.DataFrame] _edge_target : Optional[pd.DataFrame] + _edge_target_raw : Optional[pd.DataFrame] _weighted_adjacency: Optional[Any] _weighted_adjacency_nodes : Optional[Any] @@ -119,12 +135,14 @@ class Plottable(object): _dbscan_edges: Optional[Any] # fit model _adjacency : Optional[Any] - _entity_to_index : Optional[dict] - _index_to_entity : Optional[dict] + _entity_to_index : Optional[Dict] + _index_to_entity : Optional[Dict] + # DGL DGL_graph: Optional[Any] + _dgl_graph: Optional[Any] - # embed utils + # KG embeddings _relation : Optional[str] _use_feat: bool _triplets: Optional[List] # actually torch.Tensor too @@ -134,56 +152,83 @@ class Plottable(object): _partition_offsets: Optional[Dict[str, Dict[int, float]]] # from gib - def __init__(self, *args, **kwargs): - #raise RuntimeError('should not happen') - None + def reset_caches(self) -> None: + ... + + def addStyle( + self, + fg: Dict[str, Any], + bg: Optional[Dict[str, Any]] = None, + page: Optional[Dict[str, Any]] = None, + logo: Optional[Dict[str, Any]] = None, + ) -> 'Plottable': + ... + + def name(self, name: str) -> 'Plottable': + ... + + def description(self, description: str) -> 'Plottable': + ... def nodes( self, nodes: Union[Callable, Any], node: Optional[str] = None, - *args, **kwargs + *args: Any, **kwargs: Any ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def edges( self, edges: Union[Callable, Any], source: Optional[str] = None, destination: Optional[str] = None, edge: Optional[str] = None, - *args, **kwargs + *args: Any, **kwargs: Any ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self - - def pipe(self, graph_transform: Callable, *args, **kwargs) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self - - def bind(self, source=None, destination=None, node=None, edge=None, - edge_title=None, edge_label=None, edge_color=None, edge_weight=None, edge_size=None, edge_opacity=None, edge_icon=None, - edge_source_color=None, edge_destination_color=None, - point_title=None, point_label=None, point_color=None, point_weight=None, point_size=None, point_opacity=None, point_icon=None, - point_x=None, point_y=None): - if 1 + 1: - raise RuntimeError('should not happen') - return self - - def copy(self): - if 1 + 1: - raise RuntimeError('should not happen') - return self - - # ### compute + ... + + def pipe(self, graph_transform: Callable, *args: Any, **kwargs: Any) -> 'Plottable': + ... + + def graph(self, ig: Any) -> 'Plottable': + ... + + def bind( + self, + source: Optional[str] = None, + destination: Optional[str] = None, + node: Optional[str] = None, + edge: Optional[str] = None, + edge_title: Optional[str] = None, + edge_label: Optional[str] = None, + edge_color: Optional[str] = None, + edge_weight: Optional[str] = None, + edge_size: Optional[str] = None, + edge_opacity: Optional[str] = None, + edge_icon: Optional[str] = None, + edge_source_color: Optional[str] = None, + edge_destination_color: Optional[str] = None, + point_title: Optional[str] = None, + point_label: Optional[str] = None, + point_color: Optional[str] = None, + point_weight: Optional[str] = None, + point_size: Optional[str] = None, + point_opacity: Optional[str] = None, + point_icon: Optional[str] = None, + point_x: Optional[str] = None, + point_y: Optional[str] = None, + dataset_id: Optional[str] = None, + url: Optional[str] = None, + nodes_file_id: Optional[str] = None, + edges_file_id: Optional[str] = None, + ) -> 'Plottable': + ... + + def copy(self) -> 'Plottable': + ... + + # ### ComputeMixin def get_indegrees(self, col: str = 'degree_in') -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def get_outdegrees(self, col: str = 'degree_out') -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def get_degrees( self, @@ -191,14 +236,10 @@ def get_degrees( degree_in: str = "degree_in", degree_out: str = "degree_out", ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def materialize_nodes(self, reuse: bool = True, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def get_topological_levels( self, @@ -207,32 +248,16 @@ def get_topological_levels( warn_cycles: bool = True, remove_self_loops: bool = True ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... - def drop_nodes( - self, - nodes: Any - ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + def drop_nodes(self, nodes: Any) -> 'Plottable': + ... - def keep_nodes( - self, - nodes: Union[List, Any] - ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + def keep_nodes(self, nodes: Union[List, Any]) -> 'Plottable': + ... - def prune_self_edges( - self - ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + def prune_self_edges(self) -> 'Plottable': + ... def collapse( self, @@ -243,9 +268,7 @@ def collapse( unwrap: bool = False, verbose: bool = False ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def hop(self, nodes: Optional[pd.DataFrame], @@ -261,28 +284,20 @@ def hop(self, return_as_wave_front: bool = False, target_wave_front: Optional[pd.DataFrame] = None ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def filter_nodes_by_dict(self, filter_dict: Optional[dict] = None) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def filter_edges_by_dict(self, filter_dict: Optional[dict] = None) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... # FIXME python recursive typing issues def chain(self, ops: Union[Any, List[Any]]) -> 'Plottable': """ ops is Union[List[ASTObject], Chain] """ - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def chain_remote( self: 'Plottable', @@ -299,9 +314,7 @@ def chain_remote( """ chain is Union[List[ASTObject], Chain] """ - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def chain_remote_shape( self: 'Plottable', @@ -317,9 +330,7 @@ def chain_remote_shape( """ chain is Union[List[ASTObject], Chain] """ - if 1 + 1: - raise RuntimeError('should not happen') - return pd.DataFrame({}) + ... def python_remote_g( self: 'Plottable', @@ -332,9 +343,7 @@ def python_remote_g( run_label: Optional[str] = None, validate: bool = True ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def python_remote_table( self: 'Plottable', @@ -347,9 +356,7 @@ def python_remote_table( run_label: Optional[str] = None, validate: bool = True ) -> pd.DataFrame: - if 1 + 1: - raise RuntimeError('should not happen') - return pd.DataFrame({}) + ... def python_remote_json( self: 'Plottable', @@ -360,38 +367,30 @@ def python_remote_json( run_label: Optional[str] = None, validate: bool = True ) -> Any: - if 1 + 1: - raise RuntimeError('should not happen') - return {} + ... def to_igraph(self, directed: bool = True, - use_vids: bool = False, include_nodes: bool = True, node_attributes: Optional[List[str]] = None, - edge_attributes: Optional[List[str]] = None + edge_attributes: Optional[List[str]] = None, + use_vids: bool = False ) -> Any: - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def from_igraph(self, - ig, + ig: Any, node_attributes: Optional[List[str]] = None, edge_attributes: Optional[List[str]] = None, load_nodes: bool = True, load_edges: bool = True, merge_if_existing: bool = True - ): - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> 'Plottable': + ... def compute_igraph(self, alg: str, out_col: Optional[str] = None, directed: Optional[bool] = None, use_vids: bool = False, params: dict = {}, stringify_rich_types: bool = True ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def layout_igraph(self, layout: str, @@ -403,9 +402,7 @@ def layout_igraph(self, play: Optional[int] = 0, params: dict = {} ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def from_cugraph(self, G, @@ -413,10 +410,8 @@ def from_cugraph(self, edge_attributes: Optional[List[str]] = None, load_nodes: bool = True, load_edges: bool = True, merge_if_existing: bool = True - ): - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> 'Plottable': + ... def to_cugraph(self, directed: bool = True, @@ -424,19 +419,15 @@ def to_cugraph(self, node_attributes: Optional[List[str]] = None, edge_attributes: Optional[List[str]] = None, kind : CuGraphKind = 'Graph' - ) -> Any: - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> Union[Graph, MultiGraph, BiPartiteGraph]: + ... def compute_cugraph(self, alg: str, out_col: Optional[str] = None, params: dict = {}, kind : CuGraphKind = 'Graph', directed = True, G: Optional[Any] = None - ): - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> 'Plottable': + ... def layout_cugraph(self, layout: str = 'force_atlas2', params: dict = {}, @@ -446,20 +437,32 @@ def layout_cugraph(self, x_out_col: str = 'x', y_out_col: str = 'y', play: Optional[int] = 0 - ): - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> 'Plottable': + ... + + def layout_graphviz(self, + prog: Prog = 'dot', + args: Optional[str] = None, + directed: bool = True, + strict: bool = False, + graph_attr: Optional[Dict[GraphAttr, Any]] = None, + node_attr: Optional[Dict[NodeAttr, Any]] = None, + edge_attr: Optional[Dict[EdgeAttr, Any]] = None, + skip_styling: bool = False, + render_to_disk: bool = False, # unsafe in server settings + path: Optional[str] = None, + format: Optional[Format] = None + ) -> 'Plottable': + ... def from_networkx(self, G: Any) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... + + def networkx_checkoverlap(self, g: Any) -> None: + ... def networkx2pandas(self, G: Any) -> Tuple[pd.DataFrame, pd.DataFrame]: - if 1 + 1: - raise RuntimeError('should not happen') - return pd.DataFrame(), pd.DataFrame() + ... def fa2_layout( self, @@ -469,9 +472,7 @@ def fa2_layout( partition_key: Optional[str] = None, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def layout_settings( self, @@ -495,10 +496,8 @@ def layout_settings( precision_vs_speed: Optional[float] = None, gravity: Optional[float] = None, scaling_ratio: Optional[float] = None - ): - if 1 + 1: - raise RuntimeError('should not happen') - return self + ) -> 'Plottable': + ... def scene_settings( self, @@ -510,58 +509,41 @@ def scene_settings( edge_opacity: Optional[float] = None, point_opacity: Optional[float] = None, ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... - def encode_axis(self, rows=[]) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + def encode_axis(self, rows: List[Dict] = []) -> 'Plottable': + ... def settings(self, - height: Optional[float] = None, + height: Optional[int] = None, url_params: Dict[str, Any] = {}, render: Optional[Union[bool, RenderModes]] = None ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self - + ... + + def privacy(self, mode: Optional[PrivacyMode] = None, notify: Optional[bool] = None, invited_users: Optional[List[str]] = None, message: Optional[str] = None) -> 'Plottable': + ... + def to_cudf(self) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def to_pandas(self) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def protocol(self, v: Optional[str] = None) -> str: - if 1 + 1: - raise RuntimeError('should not happen') - return '' + ... def server(self, v: Optional[str] = None) -> str: - if 1 + 1: - raise RuntimeError('should not happen') - return '' + ... def client_protocol_hostname(self, v: Optional[str] = None) -> str: - if 1 + 1: - raise RuntimeError('should not happen') - return '' + ... def base_url_server(self, v: Optional[str] = None) -> str: - if 1 + 1: - raise RuntimeError('should not happen') - return '' + ... def base_url_client(self, v: Optional[str] = None) -> str: - if 1 + 1: - raise RuntimeError('should not happen') - return '' + ... def upload( self, @@ -569,20 +551,114 @@ def upload( erase_files_on_fail=True, validate: bool = True ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... def plot( self, - graph=None, - nodes=None, - name=None, - description=None, + graph: Optional[Any] = None, + nodes: Optional[Any] = None, + name: Optional[str] = None, + description: Optional[str] = None, render: Optional[Union[bool, RenderModes]] = "auto", - skip_upload=False, as_files=False, memoize=True, erase_files_on_fail: bool = True, - extra_html="", override_html_style=None, validate: bool = True + skip_upload: bool = False, + as_files: bool = False, + memoize: bool = True, + erase_files_on_fail: bool = True, + extra_html: str = "", + override_html_style: Optional[str] = None, + validate: bool = True + ) -> Any: + ... + + def pandas2igraph(self, edges: pd.DataFrame, directed: bool = True) -> Any: + ... + + def igraph2pandas(self, ig: Any) -> Tuple[pd.DataFrame, pd.DataFrame]: + ... + + def infer_labels(self) -> 'Plottable': + ... + + + @overload + def transform(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: str = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[True] = True, + scaled: bool = True, + verbose: bool = False) -> 'Plottable': + ... + + @overload + def transform(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: str = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[False] = False, + scaled: bool = True, + verbose: bool = False) -> Tuple[pd.DataFrame, pd.DataFrame]: + ... + + def transform(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: str = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: bool = True, + scaled: bool = True, + verbose: bool = False) -> Union[Tuple[pd.DataFrame, pd.DataFrame], 'Plottable']: + ... + + + @overload + def transform_umap(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: GraphEntityKind = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[True] = True, + fit_umap_embedding: bool = True, + umap_transform_kwargs: Dict[str, Any] = {} ) -> 'Plottable': - if 1 + 1: - raise RuntimeError('should not happen') - return self + ... + + @overload + def transform_umap(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: GraphEntityKind = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[False] = False, + fit_umap_embedding: bool = True, + umap_transform_kwargs: Dict[str, Any] = {} + ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + ... + + def transform_umap(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: GraphEntityKind = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: bool = True, + fit_umap_embedding: bool = True, + umap_transform_kwargs: Dict[str, Any] = {} + ) -> Union[Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], 'Plottable']: + ... + + + diff --git a/graphistry/PlotterBase.py b/graphistry/PlotterBase.py index 646600032c..4ccc735b2b 100644 --- a/graphistry/PlotterBase.py +++ b/graphistry/PlotterBase.py @@ -1,5 +1,5 @@ from graphistry.Plottable import Plottable, RenderModes, RenderModesConcrete -from typing import Any, Callable, Dict, List, Optional, Union +from typing import Any, Callable, Dict, List, Optional, Union, Tuple from graphistry.render.resolve_render_mode import resolve_render_mode import copy, hashlib, numpy as np, pandas as pd, pyarrow as pa, sys, uuid from functools import lru_cache @@ -9,21 +9,11 @@ from .constants import SRC, DST, NODE -from .plugins_types import CuGraphKind from .plugins_types.kusto_types import KustoConfig from .plugins_types.spanner_types import SpannerConfig -from .plugins.igraph import ( - to_igraph as to_igraph_base, from_igraph as from_igraph_base, - compute_igraph as compute_igraph_base, - layout_igraph as layout_igraph_base -) -from .plugins.graphviz import layout_graphviz as layout_graphviz_base -from .plugins_types.graphviz_types import EdgeAttr, Format, GraphAttr, NodeAttr, Prog -from .plugins.cugraph import ( - to_cugraph as to_cugraph_base, from_cugraph as from_cugraph_base, - compute_cugraph as compute_cugraph_base, - layout_cugraph as layout_cugraph_base -) +from .plugins.igraph import to_igraph, from_igraph, compute_igraph, layout_igraph +from .plugins.graphviz import layout_graphviz +from .plugins.cugraph import to_cugraph, from_cugraph, compute_cugraph, layout_cugraph from .util import ( error, hash_pdf, in_ipython, in_databricks, make_iframe, random_string, warn, cache_coercion, cache_coercion_helper, WeakValueWrapper @@ -129,7 +119,7 @@ def reset_caches(self): def __init__(self, *args: Any, **kwargs: Any) -> None: - super(PlotterBase, self).__init__() + super().__init__(*args, **kwargs) # Bindings self._edges : Any = None @@ -189,17 +179,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._node_features_raw = None self._node_target = None self._node_target_encoder = None + self._node_target_raw: Optional[pd.DataFrame] = None self._edge_embedding = None - self._edge_encoder = None + self._edge_encoder: Optional[Any] = None self._edge_features = None self._edge_features_raw = None - self._edge_target = None - self._edge_target_encoder = None + self._edge_target : Optional[pd.DataFrame] = None + self._edge_target_raw = None self._weighted_adjacency_nodes = None self._weighted_adjacency_edges = None - self._weighted_edges_df = None + self._weighted_edges_df: Optional[pd.DataFrame] = None self._weighted_edges_df_from_nodes = None self._weighted_edges_df_from_edges = None self._xy = None @@ -210,6 +201,16 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._umap_params : Optional[Dict[str, Any]] = None self._umap_fit_kwargs : Optional[Dict[str, Any]] = None self._umap_transform_kwargs : Optional[Dict[str, Any]] = None + + self._local_connectivity: int = 1 + self._metric: str = "euclidean" + self._suffix: str = "" + self._n_components: int = 2 + self._n_neighbors: int = 12 + self._negative_sample_rate: int = 5 + self._min_dist: float = 0.1 + self._repulsion_strength: float = 1 + self._spread: float = 0.5 self._dbscan_engine = None self._dbscan_params = None @@ -217,15 +218,18 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._dbscan_edges = None # fit model self._adjacency = None - self._entity_to_index = None - self._index_to_entity = None + self._entity_to_index: Optional[Dict] = None + self._index_to_entity: Optional[Dict] = None # KG embeddings self._relation : Optional[str] = None - self._use_feat: bool = False + self._use_feat = False self._triplets: Optional[List] = None self._kg_embed_dim: int = 128 + # Layout + self._partition_offsets: Optional[Dict[str, Dict[int, float]]] = None + # DGL self.DGL_graph = None # the DGL graph @@ -1224,7 +1228,7 @@ def fill_missing_bindings(g, source='src', destination='dst): return graph_transform(self, *args, **kwargs) - def graph(self, ig): + def graph(self, ig: Any) -> Plottable: """Specify the node and edge data. :param ig: NetworkX graph or an IGraph graph with node and edge attributes. @@ -1619,60 +1623,13 @@ def plot( else: raise ValueError(f"Unexpected render mode resolution: {render_mode}") - def from_igraph(self, - ig, - node_attributes: Optional[List[str]] = None, - edge_attributes: Optional[List[str]] = None, - load_nodes = True, load_edges = True, - merge_if_existing = True - ): - return from_igraph_base( - self, - ig, - node_attributes, - edge_attributes, - load_nodes, load_edges, - merge_if_existing - ) - from_igraph.__doc__ = from_igraph_base.__doc__ - - - def to_igraph(self, - directed = True, use_vids = False, include_nodes = True, - node_attributes: Optional[List[str]] = None, - edge_attributes: Optional[List[str]] = None, - ): - return to_igraph_base( - self, - directed = directed, use_vids = use_vids, include_nodes = include_nodes, - node_attributes = node_attributes, - edge_attributes = edge_attributes - ) - to_igraph.__doc__ = to_igraph_base.__doc__ - - - def compute_igraph(self, - alg: str, out_col: Optional[str] = None, directed: Optional[bool] = None, use_vids: bool = False, params: dict = {}, stringify_rich_types: bool = True - ): - return compute_igraph_base(self, alg, out_col, directed, use_vids, params, stringify_rich_types) - compute_igraph.__doc__ = compute_igraph_base.__doc__ - - - def layout_igraph(self, - layout: str, - directed: Optional[bool] = None, - use_vids: bool = False, - bind_position: bool = True, - x_out_col: str = 'x', - y_out_col: str = 'y', - play: Optional[int] = 0, - params: dict = {} - ): - return layout_igraph_base(self, layout, directed, use_vids, bind_position, x_out_col, y_out_col, play, params) - layout_igraph.__doc__ = layout_igraph_base.__doc__ + from_igraph = from_igraph + to_igraph = to_igraph + compute_igraph = compute_igraph + layout_igraph = layout_igraph - def pandas2igraph(self, edges, directed=True): + def pandas2igraph(self, edges: pd.DataFrame, directed: bool = True) -> Any: """Convert a pandas edge dataframe to an IGraph graph. Uses current bindings. Defaults to treating edges as directed. @@ -1708,7 +1665,7 @@ def pandas2igraph(self, edges, directed=True): vertex_name_attr=self._node) - def igraph2pandas(self, ig): + def igraph2pandas(self, ig: Any) -> Tuple[pd.DataFrame, pd.DataFrame]: """Under current bindings, transform an IGraph into a pandas edges dataframe and a nodes dataframe. Deprecated in favor of `.from_igraph()` @@ -1755,7 +1712,7 @@ def get_edgelist(ig): return (edges, nodes) - def networkx_checkoverlap(self, g): + def networkx_checkoverlap(self, g: Any) -> None: """ Raise an error if the node attribute already exists in the graph """ @@ -1771,7 +1728,7 @@ def networkx_checkoverlap(self, g): if not (self._node is None) and self._node in vattribs: error('Vertex attribute "%s" already exists.' % self._node) - def networkx2pandas(self, g): + def networkx2pandas(self, G: Any) -> Tuple[pd.DataFrame, pd.DataFrame]: def get_nodelist(g): for n in g.nodes(data=True): @@ -1782,11 +1739,11 @@ def get_edgelist(g): yield dict({self._source: e[0], self._destination: e[1]}, **e[2]) self._check_mandatory_bindings(False) - self.networkx_checkoverlap(g) + self.networkx_checkoverlap(G) self._node = self._node or PlotterBase._defaultNodeId - nodes = pd.DataFrame(get_nodelist(g)) - edges = pd.DataFrame(get_edgelist(g)) + nodes = pd.DataFrame(get_nodelist(G)) + edges = pd.DataFrame(get_edgelist(G)) return (edges, nodes) def from_networkx(self, G) -> Plottable: @@ -1870,72 +1827,13 @@ def from_networkx(self, G) -> Plottable: e_df, n_df = g.networkx2pandas(G) return g.edges(e_df).nodes(n_df) - def from_cugraph(self, - G, - node_attributes: Optional[List[str]] = None, - edge_attributes: Optional[List[str]] = None, - load_nodes: bool = True, load_edges: bool = True, - merge_if_existing: bool = True - ): - return from_cugraph_base( - self, G, - node_attributes, edge_attributes, load_nodes, merge_if_existing) - from_cugraph.__doc__ = from_cugraph_base.__doc__ - - def to_cugraph(self, - directed: bool = True, - include_nodes: bool = True, - node_attributes: Optional[List[str]] = None, - edge_attributes: Optional[List[str]] = None, - kind : CuGraphKind = 'Graph' - ): - return to_cugraph_base( - self, directed, include_nodes, node_attributes, edge_attributes, kind - ) - to_cugraph.__doc__ = to_cugraph_base.__doc__ - def compute_cugraph(self, - alg: str, out_col: Optional[str] = None, params: dict = {}, - kind : CuGraphKind = 'Graph', directed = True, - G: Optional[Any] = None - ): - return compute_cugraph_base( - self, alg, out_col, params, kind, directed, G - ) - compute_cugraph.__doc__ = compute_cugraph_base.__doc__ - - def layout_cugraph(self, - layout: str = 'force_atlas2', params: dict = {}, - kind : CuGraphKind = 'Graph', directed = True, - G: Optional[Any] = None, - bind_position: bool = True, - x_out_col: str = 'x', - y_out_col: str = 'y', - play: Optional[int] = 0 - ): - return layout_cugraph_base( - self, layout, params, kind, directed, G, - bind_position, x_out_col, y_out_col, play - ) - layout_cugraph.__doc__ = layout_cugraph_base.__doc__ - - def layout_graphviz(self, - prog: Prog = 'dot', - args: Optional[str] = None, - directed: bool = True, - strict: bool = False, - graph_attr: Optional[Dict[GraphAttr, Any]] = None, - node_attr: Optional[Dict[NodeAttr, Any]] = None, - edge_attr: Optional[Dict[EdgeAttr, Any]] = None, - skip_styling: bool = False, - render_to_disk: bool = False, # unsafe in server settings - path: Optional[str] = None, - format: Optional[Format] = None - ) -> Plottable: - return layout_graphviz_base( - self, prog, args, directed, strict, graph_attr, node_attr, edge_attr, skip_styling, render_to_disk, path, format - ) - layout_graphviz.__doc__ = layout_graphviz_base.__doc__ + from_cugraph = from_cugraph + to_cugraph = to_cugraph + compute_cugraph = compute_cugraph + layout_cugraph = layout_cugraph + + layout_graphviz = layout_graphviz def _check_mandatory_bindings(self, node_required): if self._source is None or self._destination is None: @@ -2289,7 +2187,6 @@ def bolt(self, driver): def spanner(self: Plottable, spanner_config: SpannerConfig) -> Plottable: """ Set spanner configuration for this Plottable. - SpannerConfig - "instance_id": The Spanner instance ID. - "database_id": The Spanner database ID. @@ -2306,12 +2203,11 @@ def spanner(self: Plottable, spanner_config: SpannerConfig) -> Plottable: """ self._spanner_config = spanner_config return self - + def kusto(self: Plottable, kusto_config: KustoConfig) -> Plottable: """ Set kusto configuration for this Plottable. - KustoConfig - "cluster": The Kusto cluster name. - "database": The Kusto database name. @@ -2320,7 +2216,6 @@ def kusto(self: Plottable, kusto_config: KustoConfig) -> Plottable: - "client_secret": The Kusto client secret. - "tenant_id": The Kusto tenant ID. Otherwise: process will use web browser to authenticate. - :param kusto_config: A dictionary containing the Kusto configuration. :type (KustoConfig) :returns: Plottable with a Kusto connection @@ -2328,7 +2223,7 @@ def kusto(self: Plottable, kusto_config: KustoConfig) -> Plottable: """ self._kusto_config = kusto_config return self - + def infer_labels(self): """ @@ -2518,7 +2413,7 @@ def cypher(self, query: str, params: Dict[str, Any] = {}) -> Plottable: .nodes(nodes)\ .edges(edges) - + def spanner_gql_to_g(self: Plottable, query: str) -> Plottable: """ Submit GQL query to google spanner graph database and return Plottable with nodes and edges populated @@ -2526,34 +2421,24 @@ def spanner_gql_to_g(self: Plottable, query: str) -> Plottable: GQL must be a path query with a syntax similar to the following, it's recommended to return the path with SAFE_TO_JSON(p), TO_JSON() can also be used, but not recommend. LIMIT is optional, but for large graphs with millions of edges or more, it's best to filter either in the query or use LIMIT so as not to exhaust GPU memory. - query=f'''GRAPH my_graph MATCH p = (a)-[b]->(c) LIMIT 100000 return SAFE_TO_JSON(p) as path''' - :param query: GQL query string :type query: Str - :returns: Plottable with the results of GQL query as a graph :rtype: Plottable - **Example: calling spanner_gql_to_g :: - import graphistry - # credentials_file is optional, all others are required SPANNER_CONF = { "project_id": PROJECT_ID, "instance_id": INSTANCE_ID, "database_id": DATABASE_ID, "credentials_file": CREDENTIALS_FILE } - graphistry.register(..., spanner_config=SPANNER_CONF) - query=f'''GRAPH my_graph MATCH p = (a)-[b]->(c) LIMIT 100000 return SAFE_TO_JSON(p) as path''' - g = graphistry.spanner_gql_to_g(query) - g.plot() """ @@ -2561,96 +2446,76 @@ def spanner_gql_to_g(self: Plottable, query: str) -> Plottable: res = copy.copy(self) with SpannerGraphContext(res._spanner_config) as sg: return sg.gql_to_graph(query, g=res) - + def spanner_query_to_df(self: Plottable, query: str) -> pd.DataFrame: """ - Submit query to google spanner database and return a df of the results query can be SQL or GQL as long as table of results are returned - query='SELECT * from Account limit 10000' - :param query: query string :type query: Str - :returns: Pandas DataFrame with the results of query :rtype: pd.DataFrame - **Example: calling spanner_query_to_df :: - import graphistry - # credentials_file is optional, all others are required SPANNER_CONF = { "project_id": PROJECT_ID, "instance_id": INSTANCE_ID, "database_id": DATABASE_ID, "credentials_file": CREDENTIALS_FILE } - graphistry.register(..., spanner_config=SPANNER_CONF) - query='SELECT * from Account limit 10000' - df = graphistry.spanner_query_to_df(query) - g.plot() """ from .plugins.spannergraph import SpannerGraphContext with SpannerGraphContext(self._spanner_config) as sg: return sg.query_to_df(query) - + def kusto_query(self: Plottable, query: str, unwrap_nested: Optional[bool] = None) -> List[pd.DataFrame]: """ Submit a Kusto/Azure Data Explorer *query* and return result tables. - Because a Kusto request may emit multiple tables, a **list of DataFrames** is always returned; most queries yield a single entry. - unwrap_nested ------------- Controls auto-flattening of *dynamic* (JSON) columns: • True - always try to flatten, raise if it fails • None - default heuristic: flatten only if table looks nested • False - leave results untouched - :param query: Kusto query string :type query: str :param unwrap_nested: flatten strategy above :type unwrap_nested: bool | None :returns: list of Pandas DataFrames :rtype: List[pd.DataFrame] - **Example** :: - frames = graphistry.kusto_query("StormEvents | take 100") df = frames[0] """ from .plugins.kustograph import KustoGraphContext with KustoGraphContext(self._kusto_config) as kg: return kg.query(query, unwrap_nested=unwrap_nested) - + def kusto_query_graph(self: Plottable, graph_name: str, snap_name: Optional[str] = None) -> Plottable: """ Fetch a Kusto *graph* (and optional *snapshot*) as a Graphistry object. - Under the hood: `graph(..)` + `graph-to-table` to pull **nodes** and **edges**, then binds them to *self*. - :param graph_name: name of Kusto graph entity :type graph_name: str :param snap_name: optional snapshot/version :type snap_name: str | None :returns: Plottable ready for `.plot()` or further transforms :rtype: Plottable - **Example** :: - g = graphistry.kusto_query_graph("HoneypotNetwork").plot() """ from .plugins.kustograph import KustoGraphContext diff --git a/graphistry/compute/ComputeMixin.py b/graphistry/compute/ComputeMixin.py index 8aead43a91..2b5ffd7710 100644 --- a/graphistry/compute/ComputeMixin.py +++ b/graphistry/compute/ComputeMixin.py @@ -1,5 +1,5 @@ import numpy as np, pandas as pd -from typing import Any, List, Union, TYPE_CHECKING +from typing import Any, List, Union from inspect import getmodule from graphistry.Engine import Engine, EngineAbstract @@ -21,17 +21,12 @@ logger = setup_logger(__name__) -if TYPE_CHECKING: - MIXIN_BASE = Plottable -else: - MIXIN_BASE = object - - -class ComputeMixin(MIXIN_BASE): - def __init__(self, *args, **kwargs): - pass - +class ComputeMixin(Plottable): + + def __init__(self, *a, **kw): + super().__init__(*a, **kw) + def to_cudf(self) -> 'Plottable': """ Convert to GPU mode by converting any defined nodes and edges to cudf dataframes diff --git a/graphistry/compute/cluster.py b/graphistry/compute/cluster.py index de49223ed1..8d1f5a9149 100644 --- a/graphistry/compute/cluster.py +++ b/graphistry/compute/cluster.py @@ -1,4 +1,4 @@ -from typing import Any, List, Union, TYPE_CHECKING, Tuple, Optional +from typing import Any, List, Union, TYPE_CHECKING, Tuple, Optional, cast from typing_extensions import Literal from collections import Counter from inspect import getmodule @@ -262,8 +262,9 @@ def dbscan_predict_cuml(X: Any, model: Any) -> Any: class ClusterMixin(MIXIN_BASE): - def __init__(self, *args, **kwargs): - pass + + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def _cluster_dbscan( self, kind: GraphEntityKind, cols, fit_umap_embedding, target, min_dist, min_samples, engine_dbscan: DBSCANEngineAbstract, verbose, *args, **kwargs @@ -376,11 +377,11 @@ def dbscan( return res def _transform_dbscan( - self: Plottable, df: pd.DataFrame, ydf, kind, verbose + self, df: pd.DataFrame, ydf, kind, verbose ) -> Tuple[Union[pd.DataFrame, None], pd.DataFrame, pd.DataFrame, pd.DataFrame]: res = self.bind() - if hasattr(res, "_dbscan_params"): + if hasattr(res, "_dbscan_params") and res._dbscan_params is not None: # Assume that we are transforming to last fit of dbscan cols = res._dbscan_params["cols"] umap = res._dbscan_params["fit_umap_embedding"] diff --git a/graphistry/compute/conditional.py b/graphistry/compute/conditional.py index 413e8d8724..e031ae426a 100644 --- a/graphistry/compute/conditional.py +++ b/graphistry/compute/conditional.py @@ -57,8 +57,9 @@ def probs(x, given, df: pd.DataFrame, how='index'): return res class ConditionalMixin(MIXIN_BASE): - def __init__(self, *args, **kwargs): - pass + + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def conditional_graph(self, x, given, kind='nodes', *args, **kwargs): """ diff --git a/graphistry/compute/filter_by_dict.py b/graphistry/compute/filter_by_dict.py index ea6350d086..31d17726cc 100644 --- a/graphistry/compute/filter_by_dict.py +++ b/graphistry/compute/filter_by_dict.py @@ -49,7 +49,7 @@ def filter_by_dict(df: DataFrameT, filter_dict: Optional[dict] = None, engine: U return df[hits] -def filter_nodes_by_dict(self: Plottable, filter_dict: dict, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO) -> Plottable: +def filter_nodes_by_dict(self: Plottable, filter_dict: Optional[dict] = None, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO) -> Plottable: """ filter nodes to those that match all values in filter_dict """ @@ -57,7 +57,7 @@ def filter_nodes_by_dict(self: Plottable, filter_dict: dict, engine: Union[Engin return self.nodes(nodes2) -def filter_edges_by_dict(self: Plottable, filter_dict: dict, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO) -> Plottable: +def filter_edges_by_dict(self: Plottable, filter_dict: Optional[dict] = None, engine: Union[EngineAbstract, str] = EngineAbstract.AUTO) -> Plottable: """ filter edges to those that match all values in filter_dict """ diff --git a/graphistry/dgl_utils.py b/graphistry/dgl_utils.py index a49faec41d..0afd49facf 100644 --- a/graphistry/dgl_utils.py +++ b/graphistry/dgl_utils.py @@ -1,7 +1,7 @@ # classes for converting a dataframe or Graphistry Plottable into a DGL from collections import Counter from inspect import getmodule -from typing import Dict, Optional, TYPE_CHECKING, Tuple +from typing import Dict, Optional, TYPE_CHECKING, Tuple, cast, Any import warnings import numpy as np @@ -214,12 +214,17 @@ def get_torch_train_test_mask(n: int, ratio: float = 0.8): ####################################################################################################################### -class DGLGraphMixin(MIXIN_BASE): +class DGLGraphMixin(FeatureMixin): """ Automagic DGL models from Graphistry Instances. """ - def __init__(self): + _dgl_graph: Optional[Any] + _edges: Any + _edge_target : Optional[pd.DataFrame] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) self.dgl_initialized = False def dgl_lazy_init(self, train_split: float = 0.8, device: str = "cpu"): @@ -265,8 +270,8 @@ def _remove_edges_not_in_nodes(self, node_column: str): assert ( mask.sum() > 2 ), f"mask slice is (practically) empty, will lead to bad graph, found {sum(mask)}" - self._MASK = mask # type: ignore - self._edges = edf[mask] # type: ignore + self._MASK = mask + self._edges = edf[mask] logger.debug(f'new EDGES: length: {len(self._edges)}') @@ -319,7 +324,7 @@ def _convert_edge_dataframe_to_DGL( if inplace: res = self else: - res = self.bind() + res = cast('DGLGraphMixin', self.bind()) if res._node is None: res._node = config.IMPLICIT_NODE_ID @@ -453,7 +458,7 @@ def build_gnn( if inplace: res = self else: - res = self.bind() + res = cast('DGLGraphMixin', self.bind()) res.dgl_lazy_init(train_split=train_split, device=device) diff --git a/graphistry/embed_utils.py b/graphistry/embed_utils.py index 73e70c371e..aa076102a8 100644 --- a/graphistry/embed_utils.py +++ b/graphistry/embed_utils.py @@ -1,7 +1,7 @@ import logging import numpy as np import pandas as pd -from typing import Optional, Union, Callable, List, TYPE_CHECKING, Any, Tuple +from typing import Optional, Union, Callable, List, TYPE_CHECKING, Any, Tuple, cast from graphistry.utils.lazy_import import lazy_embed_import from .PlotterBase import Plottable @@ -72,9 +72,12 @@ def RotatE(h:TT, r:TT, t:TT) -> TT: # type: ignore return -(h * r - t).norm(p=1, dim=1) # type: ignore -class HeterographEmbedModuleMixin(MIXIN_BASE): - def __init__(self): - super().__init__() +class HeterographEmbedModuleMixin(ComputeMixin): + _nodes: Any + _edges: Any + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) self._protocol = { "TransE": EmbedDistScore.TransE, @@ -86,9 +89,6 @@ def __init__(self): self._relation2id = {} self._id2node = {} self._id2relation = {} - self._relation = None - self._use_feat = False - self._kg_embed_dim = None self._kg_embeddings = None self._embed_model = None @@ -319,7 +319,7 @@ def embed( if inplace: res = self else: - res = self.bind() + res = cast('HeterographEmbedModuleMixin', self.bind()) requires_new_model = False if res._relation != relation: diff --git a/graphistry/feature_utils.py b/graphistry/feature_utils.py index 106b5a646a..fc5cda9c80 100644 --- a/graphistry/feature_utils.py +++ b/graphistry/feature_utils.py @@ -6,6 +6,7 @@ from inspect import getmodule import warnings from functools import partial +from typing import cast from typing import ( List, @@ -15,7 +16,9 @@ Optional, Tuple, TYPE_CHECKING, + overload, ) +from typing_extensions import Literal from graphistry.compute.ComputeMixin import ComputeMixin from graphistry.config import config as graphistry_config @@ -2041,7 +2044,7 @@ def get_matrix_by_column_parts(X: pd.DataFrame, column_parts: Optional[Union[lis return res -class FeatureMixin(MIXIN_BASE): +class FeatureMixin(ComputeMixin): """FeatureMixin for automatic featurization of nodes and edges DataFrames. Subclasses UMAPMixin for umap-ing of automatic features. Usage: @@ -2067,8 +2070,8 @@ class FeatureMixin(MIXIN_BASE): _feature_memoize: WeakValueDictionary = WeakValueDictionary() _feature_params: dict = {} - def __init__(self, *args, **kwargs): - pass + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def _get_feature(self, kind): kind2 = kind.replace('s', '') @@ -2115,7 +2118,7 @@ def _featurize_nodes( memoize: bool = True, verbose: bool = False, ): - res = self.copy() + res = self.copy() ndf = res._nodes node = res._node @@ -2175,7 +2178,7 @@ def _featurize_nodes( feature_engine=feature_engine, ) - res._feature_params = { + cast('FeatureMixin', res)._feature_params = { **getattr(res, "_feature_params", {}), "nodes": fkwargs, } @@ -2246,17 +2249,17 @@ def _featurize_edges( verbose: bool = False, ): - res = self.copy() + res = cast('FeatureMixin', self.copy()) edf = res._edges X_resolved = resolve_X(edf, X) y_resolved = resolve_y(edf, y) - if res._source not in X_resolved: + if res._source and res._source not in X_resolved: logger.debug("adding g._source to edge features") X_resolved = X_resolved.assign( **{res._source: res._edges[res._source]} ) - if res._destination not in X_resolved: + if res._destination and res._destination not in X_resolved: logger.debug("adding g._destination to edge features") X_resolved = X_resolved.assign( **{res._destination: res._edges[res._destination]} @@ -2357,6 +2360,32 @@ def _transform(self, encoder: str, df: pd.DataFrame, ydf: Optional[pd.DataFrame] ) raise ValueError(f"Encoder {encoder} not initialized. Call featurize() first.") + @overload + def transform(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: str = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[True] = True, + scaled: bool = True, + verbose: bool = False) -> 'Plottable': + ... + + @overload + def transform(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: str = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[False] = False, + scaled: bool = True, + verbose: bool = False) -> Tuple[pd.DataFrame, pd.DataFrame]: + ... + def transform(self, df: pd.DataFrame, y: Optional[pd.DataFrame] = None, kind: str = 'nodes', @@ -2678,7 +2707,7 @@ def featurize( if inplace: res = self else: - res = self.bind() + res = cast('FeatureMixin', self.bind()) raw_feature_engine = feature_engine feature_engine = resolve_feature_engine(feature_engine) @@ -2802,7 +2831,7 @@ def featurize_or_get_nodes_dataframe_if_X_is_None( """helper method gets node feature and target matrix if X, y are not specified. if X, y are specified will set them as `_node_target` and `_node_target` attributes """ - res = self.bind() + res = cast('FeatureMixin', self.bind()) if not reuse_if_existing: # will cause re-featurization res._node_features = None @@ -2810,12 +2839,15 @@ def featurize_or_get_nodes_dataframe_if_X_is_None( if reuse_if_existing and res._node_features is not None: logger.info('-Reusing Existing Node Featurization') - return res._node_features, res._node_target, res + # NOTE: Perhaps for these df.empty stubs, we need len(y.index) == len(X.index) + # even if no columns, downstream code could be confused on shape mismatch + node_target = res._node_target if res._node_target is not None else pd.DataFrame() + return res._node_features, node_target, res use_scaler_resolved = resolve_scaler(use_scaler, feature_engine) use_scaler_target_resolved = resolve_scaler_target(use_scaler_target, feature_engine, multilabel) - res = res._featurize_nodes( + res = cast('FeatureMixin', res)._featurize_nodes( X=X, y=y, use_scaler=use_scaler_resolved, @@ -2850,7 +2882,7 @@ def featurize_or_get_nodes_dataframe_if_X_is_None( assert res._node_features is not None # ensure no infinite loop - return res.featurize_or_get_nodes_dataframe_if_X_is_None( + return cast('FeatureMixin', res).featurize_or_get_nodes_dataframe_if_X_is_None( res._node_features, res._node_target, use_scaler=use_scaler_resolved, @@ -2924,7 +2956,7 @@ def featurize_or_get_edges_dataframe_if_X_is_None( :return: data `X` and `y` """ - res = self.bind() + res = cast('FeatureMixin', self.bind()) if not reuse_if_existing: res._edge_features = None diff --git a/graphistry/gremlin.py b/graphistry/gremlin.py index 06fa614196..d1bddd9f1a 100644 --- a/graphistry/gremlin.py +++ b/graphistry/gremlin.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Iterable, List, Optional, Set, Union, TYPE_CHECKING +from typing import Any, Callable, Iterable, List, Optional, Set, Union, TYPE_CHECKING, cast import os, pandas as pd from .Plottable import Plottable try: @@ -381,17 +381,21 @@ def resultset_to_g_structured_item( DROP_QUERY = 'g.V().drop()' -class GremlinMixin(MIXIN_BASE): +class GremlinMixin(Plottable): """ Universal Gremlin<>pandas/graphistry functionality across Gremlin connectors Currently serializes queries as strings instead of bytecode in order to support cosmosdb """ - _reconnect_gremlin : Optional[Callable[['GremlinMixin'], 'GremlinMixin']] = None + _reconnect_gremlin : Optional[Callable[[Plottable], Plottable]] = None _gremlin_client : Optional[Client] def __init__(self, *args, gremlin_client: Optional[Client] = None, **kwargs): + # NOTE: Cooperative constructor passes along all args/kwargs for other mixins + kwargs['gremlin_client'] = gremlin_client + super().__init__(*args, **kwargs) + if gremlin_client is not None: self._gremlin_client = gremlin_client @@ -437,7 +441,7 @@ def gremlin_client( self._gremlin_client = gremlin_client return self - def connect(self): + def connect(self) -> Plottable: """ Use previously provided credentials to connect. Disconnect any preexisting clients. """ @@ -668,7 +672,7 @@ def fetch_nodes(self, batch_size = 1000, dry_run=False, verbose=False, ignore_er node_id = g._node if node_id is None: node_id = 'id' - g = g.bind(node=node_id) + g = cast('GremlinMixin', g.bind(node=node_id)) if nodes_df is None: edges_df = g._edges if g._edges is None: @@ -739,18 +743,11 @@ def fetch_edges(self, batch_size = 1000, dry_run=False, verbose=False, ignore_er return g2 -if TYPE_CHECKING: - COSMOS_BASE = GremlinMixin - NEPTUNE_BASE = GremlinMixin -else: - COSMOS_BASE = object - NEPTUNE_BASE = object - -class NeptuneMixin(NEPTUNE_BASE): +class NeptuneMixin(GremlinMixin): - def __init__(self, *args, **kwargs): - pass + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def neptune( self, @@ -864,16 +861,16 @@ def connect(self: NeptuneMixin) -> NeptuneMixin: return self - self._reconnect_gremlin : Optional[Callable[[NeptuneMixin], NeptuneMixin]] = connect # type: ignore + self._reconnect_gremlin = connect # type: ignore return self -class CosmosMixin(COSMOS_BASE): +class CosmosMixin(GremlinMixin): - def __init__(self, *args, **kwargs): - pass + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def cosmos( self, @@ -925,7 +922,7 @@ def connect(self: CosmosMixin) -> CosmosMixin: message_serializer=GraphSONSerializersV2d0()) return self - self._reconnect_gremlin : Optional[Callable[[CosmosMixin], CosmosMixin]] = connect # type: ignore + self._reconnect_gremlin = connect # type: ignore if gremlin_client is None: if self._reconnect_gremlin is None: diff --git a/graphistry/layout/fa2.py b/graphistry/layout/fa2.py index 1f898e45a7..6a882a85ff 100644 --- a/graphistry/layout/fa2.py +++ b/graphistry/layout/fa2.py @@ -66,7 +66,7 @@ def compute_bounding_boxes(self: Plottable, partition_key: str, engine: Engine) def fa2_layout( - g: Plottable, + self: Plottable, fa2_params: Optional[Dict[str, Any]] = None, circle_layout_params: Optional[Dict[str, Any]] = None, singleton_layout: Optional[Callable[[Plottable, Union[Tuple[float, float, float, float], Any]], Plottable]] = None, @@ -98,6 +98,7 @@ def fa2_layout( :returns: A graph object with FA2 and circle layouts applied. :rtype: graphistry.Plottable.Plottable """ + g = self if isinstance(engine, str): engine = EngineAbstract(engine) diff --git a/graphistry/layouts.py b/graphistry/layouts.py index a958a04d57..a4fdcf89cd 100644 --- a/graphistry/layouts.py +++ b/graphistry/layouts.py @@ -4,7 +4,7 @@ from .layout import ( SugiyamaLayout, circle_layout as circle_layout_base, - fa2_layout as fa2_layout_base, + fa2_layout, group_in_a_box_layout as group_in_a_box_layout_base, modularity_weighted_layout as modularity_weighted_layout_base, ring_categorical as ring_categorical_base, @@ -15,24 +15,17 @@ from .util import deprecated, setup_logger logger = setup_logger(__name__) -if TYPE_CHECKING: - MIXIN_BASE = Plottable -else: - MIXIN_BASE = object +class LayoutsMixin(Plottable): -class LayoutsMixin(MIXIN_BASE): - - def __init__(self, *args, **kwargs): - pass + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def circle_layout(self, *args, **kwargs): return circle_layout_base(self, *args, **kwargs) circle_layout.__doc__ = circle_layout_base.__doc__ - def fa2_layout(self, *args, **kwargs): - return fa2_layout_base(self, *args, **kwargs) - fa2_layout.__doc__ = fa2_layout_base.__doc__ + fa2_layout = fa2_layout def group_in_a_box_layout(self, *args, **kwargs): return group_in_a_box_layout_base(self, *args, **kwargs) diff --git a/graphistry/plotter.py b/graphistry/plotter.py index 740865b6d8..52734ab13e 100644 --- a/graphistry/plotter.py +++ b/graphistry/plotter.py @@ -1,31 +1,29 @@ from .PlotterBase import PlotterBase from .compute.ComputeMixin import ComputeMixin -from .gremlin import GremlinMixin, CosmosMixin, NeptuneMixin +from .gremlin import CosmosMixin, NeptuneMixin from .layouts import LayoutsMixin -from .feature_utils import FeatureMixin # type: ignore -from .dgl_utils import DGLGraphMixin # type: ignore -from .umap_utils import UMAPMixin # type: ignore -from .embed_utils import HeterographEmbedModuleMixin # type: ignore -from .text_utils import SearchToGraphMixin # type: ignore -from .compute.conditional import ConditionalMixin # type: ignore -from .compute.cluster import ClusterMixin # type: ignore - - -mixins = ([ - CosmosMixin, NeptuneMixin, GremlinMixin, +from .feature_utils import FeatureMixin +from .dgl_utils import DGLGraphMixin +from .umap_utils import UMAPMixin +from .embed_utils import HeterographEmbedModuleMixin +from .text_utils import SearchToGraphMixin +from .compute.conditional import ConditionalMixin +from .compute.cluster import ClusterMixin + + +# NOTE: Cooperative mixins must call: +# super().__init__(*a, **kw) in their __init__ method +# to pass along args/kwargs to the next mixin in the chain +class Plotter( + CosmosMixin, NeptuneMixin, HeterographEmbedModuleMixin, SearchToGraphMixin, DGLGraphMixin, ClusterMixin, UMAPMixin, FeatureMixin, ConditionalMixin, LayoutsMixin, - ComputeMixin, PlotterBase, object -]) - - -class Plotter( # type: ignore - *mixins # type: ignore -): # type: ignore + ComputeMixin, PlotterBase +): """ Main Plotter class for Graphistry. @@ -51,17 +49,3 @@ class Plotter( # type: ignore All attributes are inherited from the mixins and base classes. """ - def __init__(self, *args, **kwargs): - PlotterBase.__init__(self, *args, **kwargs) - ComputeMixin.__init__(self, *args, **kwargs) - LayoutsMixin.__init__(self, *args, **kwargs) - ConditionalMixin.__init__(self, *args, **kwargs) - FeatureMixin.__init__(self, *args, **kwargs) - UMAPMixin.__init__(self, *args, **kwargs) - ClusterMixin.__init__(self, *args, **kwargs) - DGLGraphMixin.__init__(self, *args, **kwargs) - SearchToGraphMixin.__init__(self, *args, **kwargs) - HeterographEmbedModuleMixin.__init__(self, *args, **kwargs) - GremlinMixin.__init__(self, *args, **kwargs) - CosmosMixin.__init__(self, *args, **kwargs) - NeptuneMixin.__init__(self, *args, **kwargs) diff --git a/graphistry/plugins/igraph.py b/graphistry/plugins/igraph.py index aa1632fc63..fadea79437 100644 --- a/graphistry/plugins/igraph.py +++ b/graphistry/plugins/igraph.py @@ -1,5 +1,5 @@ import pandas as pd -from typing import List, Optional +from typing import List, Optional, Any from graphistry.constants import NODE from graphistry.Plottable import Plottable from graphistry.util import setup_logger @@ -15,7 +15,7 @@ NODE_IGRAPH = NODE def from_igraph(self, - ig, + ig: Any, node_attributes: Optional[List[str]] = None, edge_attributes: Optional[List[str]] = None, load_nodes: bool = True, load_edges: bool = True, @@ -221,7 +221,7 @@ def to_igraph( node_attributes: Optional[List[str]] = None, edge_attributes: Optional[List[str]] = None, use_vids: bool = False -): +) -> Any: """Convert current item to igraph Graph . See examples in from_igraph. :param directed: Whether to create a directed graph (default True) diff --git a/graphistry/plugins/spannergraph.py b/graphistry/plugins/spannergraph.py index aacee9e572..e76fc149fb 100644 --- a/graphistry/plugins/spannergraph.py +++ b/graphistry/plugins/spannergraph.py @@ -259,9 +259,9 @@ def __init__(self, config: Optional[SpannerConfig] = None): if not config: raise ValueError("Missing spanner_config. Register globally with spanner_config or use with_spanner().") self.spanner_graph = SpannerGraph.from_config(config) - + def __enter__(self): return self.spanner_graph - + def __exit__(self, exc_type, exc_value, traceback): self.spanner_graph.close() diff --git a/graphistry/pygraphistry.py b/graphistry/pygraphistry.py index 921305dbd8..46ff0ac874 100644 --- a/graphistry/pygraphistry.py +++ b/graphistry/pygraphistry.py @@ -573,12 +573,11 @@ def set_bolt_driver(driver=None): PyGraphistry._config["bolt_driver"] = bolt_util.to_bolt_driver(driver) @staticmethod - # def set_spanner_config(spanner_config): - def set_spanner_config(spanner_config: Optional[Union[Dict, str]] = None): + def set_spanner_config(spanner_config: Optional[SpannerConfig] = None): """ Saves the spanner config to internal Pygraphistry _config :param spanner_config: dict of the project_id, instance_id and database_id - :type spanner_config: Optional[Union[Dict, Any]] + :type spanner_config: Optional[SpannerConfig] :returns: None. :rtype: None @@ -608,14 +607,11 @@ def set_spanner_config(spanner_config: Optional[Union[Dict, str]] = None): graphistry.set_spanner_config(SPANNER_CONF) """ - - if spanner_config is not None: - PyGraphistry._config["spanner"] = spanner_config + PyGraphistry._config["spanner"] = spanner_config @staticmethod - def set_kusto_config(kusto_config): - if kusto_config is not None: - PyGraphistry._config["kusto"] = kusto_config + def set_kusto_config(kusto_config: Optional[KustoConfig]): + PyGraphistry._config["kusto"] = kusto_config @staticmethod @@ -631,8 +627,8 @@ def register( api: Optional[Literal[1, 3]] = None, certificate_validation: Optional[bool] = None, bolt: Optional[Union[Dict, Any]] = None, - spanner_config: Optional[Union[Dict, Any]] = None, - kusto_config: Optional[Union[Dict, Any]] = None, + spanner_config: Optional[SpannerConfig] = None, + kusto_config: Optional[KustoConfig] = None, token_refresh_ms: int = 10 * 60 * 1000, store_token_creds_in_memory: Optional[bool] = None, client_protocol_hostname: Optional[str] = None, @@ -1084,11 +1080,29 @@ def bolt(driver=None): """ return Plotter().bolt(driver) + + @staticmethod + def cypher(query, params={}): + """ + + :param query: a cypher query + :param params: cypher query arguments + :return: Plotter with data from a cypher query. This call binds `source`, `destination`, and `node`. + + Call this to immediately execute a cypher query and store the graph in the resulting Plotter. + + :: + + import graphistry + g = graphistry.bolt({ query='MATCH (a)-[r:PAYMENT]->(b) WHERE r.USD > 7000 AND r.USD < 10000 RETURN r ORDER BY r.USD DESC', params={ "AccountId": 10 }) + """ + return Plotter().cypher(query, params) + + @staticmethod def spanner(spanner_config: SpannerConfig) -> Plottable: """ Set spanner configuration for this Plottable. - SpannerConfig - "instance_id": The Spanner instance ID. - "database_id": The Spanner database ID. @@ -1109,7 +1123,6 @@ def spanner(spanner_config: SpannerConfig) -> Plottable: def kusto(kusto_config: KustoConfig) -> Plottable: """ Set kusto configuration for this Plottable. - KustoConfig - "cluster": The Kusto cluster name. - "database": The Kusto database name. @@ -1118,7 +1131,6 @@ def kusto(kusto_config: KustoConfig) -> Plottable: - "client_secret": The Kusto client secret. - "tenant_id": The Kusto tenant ID. Otherwise: process will use web browser to authenticate. - :param kusto_config: A dictionary containing the Kusto configuration. :type (KustoConfig) :returns: Plottable with a Kusto connection @@ -1126,22 +1138,6 @@ def kusto(kusto_config: KustoConfig) -> Plottable: """ return Plotter().kusto(kusto_config) - @staticmethod - def cypher(query, params={}): - """ - - :param query: a cypher query - :param params: cypher query arguments - :return: Plotter with data from a cypher query. This call binds `source`, `destination`, and `node`. - - Call this to immediately execute a cypher query and store the graph in the resulting Plotter. - - :: - - import graphistry - g = graphistry.bolt({ query='MATCH (a)-[r:PAYMENT]->(b) WHERE r.USD > 7000 AND r.USD < 10000 RETURN r ORDER BY r.USD DESC', params={ "AccountId": 10 }) - """ - return Plotter().cypher(query, params) @staticmethod def nodexl(xls_or_url, source="default", engine=None, verbose=False): @@ -1919,7 +1915,7 @@ def tigergraph( return Plotter().tigergraph( protocol, server, web_port, api_port, db, user, pwd, verbose ) - + @staticmethod def spanner_gql_to_g(query: str) -> Plottable: """ @@ -1928,34 +1924,24 @@ def spanner_gql_to_g(query: str) -> Plottable: GQL must be a path query with a syntax similar to the following, it's recommended to return the path with SAFE_TO_JSON(p), TO_JSON() can also be used, but not recommend. LIMIT is optional, but for large graphs with millions of edges or more, it's best to filter either in the query or use LIMIT so as not to exhaust GPU memory. - query=f'''GRAPH my_graph MATCH p = (a)-[b]->(c) LIMIT 100000 return SAFE_TO_JSON(p) as path''' - :param query: GQL query string :type query: Str - :returns: Plottable with the results of GQL query as a graph :rtype: Plottable - **Example: calling spanner_gql_to_g :: - import graphistry - # credentials_file is optional, all others are required SPANNER_CONF = { "project_id": PROJECT_ID, "instance_id": INSTANCE_ID, "database_id": DATABASE_ID, "credentials_file": CREDENTIALS_FILE } - graphistry.register(..., spanner_config=SPANNER_CONF) - query=f'''GRAPH my_graph MATCH p = (a)-[b]->(c) LIMIT 100000 return SAFE_TO_JSON(p) as path''' - g = graphistry.spanner_gql_to_g(query) - g.plot() """ @@ -1964,36 +1950,25 @@ def spanner_gql_to_g(query: str) -> Plottable: @staticmethod def spanner_query_to_df(query: str) -> pd.DataFrame: """ - Submit query to google spanner database and return a df of the results query can be SQL or GQL as long as table of results are returned - query='SELECT * from Account limit 10000' - :param query: query string :type query: Str - :returns: Pandas DataFrame with the results of query :rtype: pd.DataFrame - **Example: calling spanner_query_to_df :: - import graphistry - # credentials_file is optional, all others are required SPANNER_CONF = { "project_id": PROJECT_ID, "instance_id": INSTANCE_ID, "database_id": DATABASE_ID, "credentials_file": CREDENTIALS_FILE } - graphistry.register(..., spanner_config=SPANNER_CONF) - query='SELECT * from Account limit 10000' - df = graphistry.spanner_query_to_df(query) - g.plot() """ @@ -2003,27 +1978,22 @@ def spanner_query_to_df(query: str) -> pd.DataFrame: def kusto_query(query: str, unwrap_nested: Optional[bool] = None) -> List[pd.DataFrame]: """ Submit a Kusto/Azure Data Explorer *query* and return result tables. - Because a Kusto request may emit multiple tables, a **list of DataFrames** is always returned; most queries yield a single entry. - unwrap_nested ------------- Controls auto‑flattening of *dynamic* (JSON) columns: • True  – always try to flatten, raise if it fails • None  – default heuristic: flatten only if table looks nested • False – leave results untouched - :param query: Kusto query string :type query: str :param unwrap_nested: flatten strategy above :type unwrap_nested: bool | None :returns: list of Pandas DataFrames :rtype: List[pd.DataFrame] - **Example** :: - frames = graphistry.kusto_query("StormEvents | take 100") df = frames[0] """ @@ -2033,24 +2003,21 @@ def kusto_query(query: str, unwrap_nested: Optional[bool] = None) -> List[pd.Dat def kusto_query_graph(graph_name: str, snap_name: Optional[str] = None) -> Plottable: """ Fetch a Kusto *graph* (and optional *snapshot*) as a Graphistry object. - Under the hood: `graph(..)` + `graph-to-table` to pull **nodes** and **edges**, then binds them to *self*. - :param graph_name: name of Kusto graph entity :type graph_name: str :param snap_name: optional snapshot/version :type snap_name: str | None :returns: Plottable ready for `.plot()` or further transforms :rtype: Plottable - **Example** :: - g = graphistry.kusto_query_graph("HoneypotNetwork").plot() """ return Plotter().kusto_query_graph(graph_name, snap_name) + @staticmethod def gsql_endpoint( self, method_name, args={}, bindings=None, db=None, dry_run=False @@ -2315,7 +2282,6 @@ def from_igraph(ig, load_nodes = True, load_edges = True ): return Plotter().from_igraph(ig, node_attributes, edge_attributes, load_nodes, load_edges) - from_igraph.__doc__ = Plotter.from_igraph.__doc__ @staticmethod def from_cugraph( @@ -2326,7 +2292,6 @@ def from_cugraph( merge_if_existing: bool = True ): return Plotter().from_cugraph(G, node_attributes, edge_attributes, load_nodes, load_edges, merge_if_existing) - from_cugraph.__doc__ = Plotter.from_cugraph.__doc__ @staticmethod def settings(height=None, url_params={}, render=None): @@ -2614,7 +2579,6 @@ def scene_settings( edge_opacity, point_opacity ) - scene_settings.__doc__ = Plotter().scene_settings.__doc__ @staticmethod diff --git a/graphistry/text_utils.py b/graphistry/text_utils.py index 5915a53014..3b44289d93 100644 --- a/graphistry/text_utils.py +++ b/graphistry/text_utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from inspect import getmodule from logging import getLogger import pandas as pd @@ -17,8 +17,8 @@ class SearchToGraphMixin(MIXIN_BASE): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **kwargs) + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def assert_fitted(self): # assert self._umap is not None, 'Umap needs to be fit first, run g.umap(..) to fit a model' @@ -213,7 +213,7 @@ def search_graph( if inplace: res = self else: - res = self.bind() + res = cast('SearchToGraphMixin', self.bind()) edf = edges = res._edges # print('shape of edges', edf.shape) diff --git a/graphistry/umap_utils.py b/graphistry/umap_utils.py index 5420b5d421..654ca3044b 100644 --- a/graphistry/umap_utils.py +++ b/graphistry/umap_utils.py @@ -1,6 +1,7 @@ import copy from time import time -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union, cast, overload +from typing_extensions import Literal from inspect import getmodule import warnings @@ -19,7 +20,7 @@ from .constants import CUML, UMAP_LEARN from .feature_utils import (FeatureMixin, XSymbolic, YSymbolic, resolve_feature_engine) -from .PlotterBase import Plottable, WeakValueDictionary +from .PlotterBase import Plottable, WeakValueDictionary, PlotterBase from .util import check_set_memoize, setup_logger @@ -152,7 +153,7 @@ def reuse_umap(g: Plottable, memoize: bool, metadata: Any) -> Optional[Plottable if o is False: return None - if isinstance(o, Plottable): + if isinstance(o, PlotterBase): return o raise ValueError(f'Expected Plottable or False, got {type(o)}') @@ -222,12 +223,9 @@ class UMAPMixin(MIXIN_BASE): """ UMAP Mixin for automagic UMAPing """ - # FIXME where is this used? - _umap_memoize: WeakValueDictionary = WeakValueDictionary() - def __init__(self, *args, **kwargs): - #self._umap_initialized = False - pass + def __init__(self, *a, **kw): + super().__init__(*a, **kw) def umap_lazy_init( @@ -387,6 +385,34 @@ def _umap_fit_transform( return emb + @overload + def transform_umap(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: GraphEntityKind = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[True] = True, + fit_umap_embedding: bool = True, + umap_transform_kwargs: Dict[str, Any] = {} + ) -> 'Plottable': + ... + + @overload + def transform_umap(self, df: pd.DataFrame, + y: Optional[pd.DataFrame] = None, + kind: GraphEntityKind = 'nodes', + min_dist: Union[str, float, int] = 'auto', + n_neighbors: int = 7, + merge_policy: bool = False, + sample: Optional[int] = None, + return_graph: Literal[False] = False, + fit_umap_embedding: bool = True, + umap_transform_kwargs: Dict[str, Any] = {} + ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: + ... + def transform_umap(self, df: pd.DataFrame, y: Optional[pd.DataFrame] = None, kind: GraphEntityKind = 'nodes', @@ -397,7 +423,7 @@ def transform_umap(self, df: pd.DataFrame, return_graph: bool = True, fit_umap_embedding: bool = True, umap_transform_kwargs: Dict[str, Any] = {} - ) -> Union[Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], Plottable]: + ) -> Union[Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame], 'Plottable']: """Transforms data into UMAP embedding Args: @@ -690,7 +716,7 @@ def umap( if inplace: res = self else: - res = self.bind() + res = cast('UMAPMixin', self.bind()) res = res.umap_lazy_init( res, @@ -899,7 +925,7 @@ def filter_weighted_edges( """ if inplace: - res = self + res: Plottable = self else: res = self.bind() diff --git a/py.typed b/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pyrightconfig.json b/pyrightconfig.json index 4548920cc0..f22e6a2ea2 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -4,5 +4,6 @@ "reportArgumentType": "none", "reportMissingTypeArgument": "none", "reportUnknownMemberType": "none", + "reportIncompatibleMethodOverride": "error", "exclude": [] }