Get pygraphistry as a singleton, remove static methods

Author: exrhizo

graphistry/PlotterBase.py

@@ -4,9 +4,10 @@
 import copy, hashlib, numpy as np, pandas as pd, pyarrow as pa, sys, uuid
 from functools import lru_cache
 from weakref import WeakValueDictionary
+import warnings
 
 from graphistry.privacy import Privacy, Mode
-from graphistry.client_session import ClientSession
+from graphistry.client_session import ClientSession, SessionManagerProtocol, DatasetInfo
 
 from .constants import SRC, DST, NODE
 from .plugins_types.kusto_types import KustoConfig
@@ -121,8 +122,17 @@ def reset_caches(self):
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
 
-        from .pygraphistry import PyGraphistry
-        self._session = kwargs.get('graphistry_client_session', PyGraphistry._session)
+        # Use late import to avoid circular dependency
+        pygraphistry = kwargs.get('pygraphistry_session', None)
+        if pygraphistry is None:
+            from .pygraphistry import PyGraphistry
+            pygraphistry = PyGraphistry
+
+            if self is PyGraphistry:
+                # NOTE: This may use global session when that isn't desired.
+                warnings.warn("Plotter initialized without pygraphistry_session, falling back to global PyGraphistry", UserWarning)
+        self._pygraphistry: SessionManagerProtocol = pygraphistry
+        self._session: ClientSession = self._pygraphistry._session
 
         # Bindings
         self._edges : Any = None
@@ -1502,21 +1512,19 @@ def server(self, v: Optional[str] = None) -> str:
 
         Note that sets are global as PyGraphistry._config entries, so be careful in multi-user environments.
         """
-        from .pygraphistry import PyGraphistry
         if v is not None:
-            PyGraphistry._session.hostname = v
-        return PyGraphistry._session.hostname
+            self._session.hostname = v
+        return self._session.hostname
 
     def protocol(self, v: Optional[str] = None) -> str:
         """
         Get or set the server protocol, e.g., "https"
 
         Note that sets are global as PyGraphistry._config entries, so be careful in multi-user environments.
         """
-        from .pygraphistry import PyGraphistry
         if v is not None:
-            PyGraphistry._session.protocol = v
-        return PyGraphistry._session.protocol
+            self._session.protocol = v
+        return self._session.protocol
 
     def client_protocol_hostname(self, v: Optional[str] = None) -> str:
         """
@@ -1526,18 +1534,15 @@ def client_protocol_hostname(self, v: Optional[str] = None) -> str:
 
         Note that sets are global as PyGraphistry._config entries, so be careful in multi-user environments.        
         """
-        from .pygraphistry import PyGraphistry
         if v is not None:
-            PyGraphistry._session.client_protocol_hostname = v
-        return PyGraphistry._session.client_protocol_hostname or f"{PyGraphistry.protocol()}://{PyGraphistry.server()}"
+            self._session.client_protocol_hostname = v
+        return self._session.client_protocol_hostname or f"{self.protocol()}://{self.server()}"
 
     def base_url_server(self, v: Optional[str] = None) -> str:
-        from .pygraphistry import PyGraphistry
-        return "%s://%s" % (PyGraphistry.protocol(), PyGraphistry.server())
+        return "%s://%s" % (self.protocol(), self.server())
 
     def base_url_client(self, v: Optional[str] = None) -> str:
-        from .pygraphistry import PyGraphistry
-        return PyGraphistry.client_protocol_hostname()
+        return self.client_protocol_hostname()
 
     def upload(
         self,
@@ -1658,8 +1663,7 @@ def plot(
                     .plot(es)
 
         """
-        from .pygraphistry import PyGraphistry
-        logger.debug("1. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
+        logger.debug("1. @PloatterBase plot: _pygraphistry.org_name: {}".format(self._session.org_name))
 
         if graph is None:
             if self._edges is None:
@@ -1673,36 +1677,34 @@ def plot(
 
         self._check_mandatory_bindings(not isinstance(n, type(None)))
 
-        # from .pygraphistry import PyGraphistry
-        api_version = PyGraphistry.api_version()
-        logger.debug("2. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
+        logger.debug("2. @PloatterBase plot: self._pygraphistry.org_name: {}".format(self._session.org_name))
         dataset: Union[ArrowUploader, Dict[str, Any], None] = None
-        if api_version == 1:
+        if self._session.api_version == 1:
             dataset = self._plot_dispatch(g, n, name, description, 'json', self._style, memoize)
             if skip_upload:
                 return dataset
-            info = PyGraphistry._etl1(dataset)
-        elif api_version == 3:
-            logger.debug("3. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
-            PyGraphistry.refresh()
-            logger.debug("4. @PloatterBase plot: PyGraphistry.org_name(): {}".format(PyGraphistry.org_name()))
+            info = self._pygraphistry._etl1(dataset)
+        elif self._session.api_version == 3:
+            logger.debug("3. @PloatterBase plot: self._pygraphistry.org_name: {}".format(self._session.org_name))
+            self._pygraphistry.refresh()
+            logger.debug("4. @PloatterBase plot: self._pygraphistry.org_name: {}".format(self._session.org_name))
 
             uploader = dataset = self._plot_dispatch_arrow(g, n, name, description, self._style, memoize)
             assert uploader is not None
             if skip_upload:
                 return uploader
-            uploader.token = PyGraphistry.api_token() # type: ignore[assignment]
+            uploader.token = self._session.api_token # type: ignore[assignment]
             uploader.post(as_files=as_files, memoize=memoize, validate=validate, erase_files_on_fail=erase_files_on_fail)
             uploader.maybe_post_share_link(self)
-            info = {
+            info: DatasetInfo = {
                 'name': uploader.dataset_id,
                 'type': 'arrow',
                 'viztoken': str(uuid.uuid4())
             }
 
-        viz_url = PyGraphistry._viz_url(info, self._url_params)
-        cfg_client_protocol_hostname = PyGraphistry._session.client_protocol_hostname
-        full_url = ('%s:%s' % (PyGraphistry._session.protocol, viz_url)) if cfg_client_protocol_hostname is None else viz_url
+        viz_url = self._pygraphistry._viz_url(info, self._url_params)
+        cfg_client_protocol_hostname = self._session.client_protocol_hostname
+        full_url = ('%s:%s' % (self._session.protocol, viz_url)) if cfg_client_protocol_hostname is None else viz_url
 
         render_mode = resolve_render_mode(self, render)
         if render_mode == "url":
@@ -2245,8 +2247,6 @@ def _make_dataset(self, edges, nodes, name, description, mode, metadata=None, me
     # Main helper for creating ETL1 payload
     def _make_json_dataset(self, edges, nodes, name) -> Dict[str, Any]:
 
-        from .pygraphistry import PyGraphistry
-
         def flatten_categorical(df):
             # Avoid cat_col.where(...)-related exceptions
             df2 = df.copy()
@@ -2260,7 +2260,7 @@ def flatten_categorical(df):
 
         bindings = {'idField': self._node or PlotterBase._defaultNodeId,
                     'destinationField': self._destination, 'sourceField': self._source}
-        dataset: Dict[str, Any] = {'name': PyGraphistry._session.dataset_prefix + name,
+        dataset: Dict[str, Any] = {'name': self._session.dataset_prefix + name,
                    'bindings': bindings, 'type': 'edgelist', 'graph': edict}
 
         if nlist is not None:
@@ -2271,20 +2271,20 @@ def flatten_categorical(df):
 
     def _make_arrow_dataset(self, edges: pa.Table, nodes: pa.Table, name: str, description: str, metadata: Optional[Dict[str, Any]]) -> ArrowUploader:
 
-        from .pygraphistry import PyGraphistry
         au : ArrowUploader = ArrowUploader(
-            server_base_path=PyGraphistry.protocol() + '://' + PyGraphistry.server(),
+            client_session=self._session,
+            server_base_path=self._session.protocol + '://' + self._session.hostname,
             edges=edges, nodes=nodes,
             name=name, description=description,
             metadata={
-                'usertag': PyGraphistry._tag,
-                'key': PyGraphistry.api_key(),
+                'usertag': self._session._tag,
+                'key': self._session.api_key,
                 'agent': 'pygraphistry',
                 'apiversion' : '3',
                 'agentversion': sys.modules['graphistry'].__version__,  # type: ignore
                 **(metadata or {})
             },
-            certificate_validation=PyGraphistry.certificate_validation())
+            certificate_validation=self._pygraphistry.certificate_validation())
 
         au.edge_encodings = au.g_to_edge_encodings(self)
         au.node_encodings = au.g_to_node_encodings(self)
@@ -2485,10 +2485,8 @@ def cypher(self, query: str, params: Dict[str, Any] = {}) -> Plottable:
         
         """
 
-        from .pygraphistry import PyGraphistry
-
         res = copy.copy(self)
-        driver = self._bolt_driver or PyGraphistry._session.bolt_driver
+        driver = self._bolt_driver or self._session.bolt_driver
         if driver is None:
             raise ValueError("BOLT connection information not provided. Must first call graphistry.register(bolt=...) or g.bolt(...).")
         with driver.session() as session:

graphistry/__init__.py

@@ -14,6 +14,7 @@
     name,
     description,
     bind,
+    client,
     style,
     addStyle,
     edges,

graphistry/client_session.py

@@ -1,11 +1,13 @@
 import os
-from typing import Any, Optional, Type, TypeVar, Union, overload, Literal, cast
+from typing import Any, Optional, Type, TypeVar, Union, overload, Literal, cast, Protocol, TypedDict, Dict
 from typing_extensions import deprecated
 from functools import lru_cache
 import json
 
 from graphistry.privacy import Privacy
 from . import util
+from .plugins_types.spanner_types import SpannerSession
+from .plugins_types.kusto_types import KustoSession
 
 
 
@@ -22,6 +24,9 @@
 class ClientSession:
 
     def __init__(self) -> None:
+        self._is_authenticated: bool = False
+        self._tag = util.fingerprint()
+
         self.api_key: Optional[str] = get_from_env(ENV_GRAPHISTRY_API_KEY, str)
         self.api_token: Optional[str] = get_from_env("GRAPHISTRY_API_TOKEN", str)
         # self.api_token_refresh_ms: Optional[int] = None
@@ -55,6 +60,33 @@ def __init__(self) -> None:
 
         # TODO: Migrate to a pattern like Kusto or Spanner
         self.bolt_driver: Optional[Any] = None
+        
+        # Plugin sessions
+        self.kusto: KustoSession = KustoSession()
+        self.spanner: SpannerSession = SpannerSession()
+
+
+class DatasetInfo(TypedDict):
+    name: str
+    viztoken: str
+    type: Literal["arrow", "vgraph"]
+
+
+
+class SessionManagerProtocol(Protocol):
+    _session: ClientSession
+
+    def _etl1(self, dataset: Any) -> DatasetInfo:
+        ...
+
+    def refresh(self, token: Optional[str] = None, fail_silent: bool = False) -> Optional[str]:
+        ...
+    
+    def _viz_url(self, info: DatasetInfo, url_params: Dict[str, Any]) -> str:
+        ...
+
+    def certificate_validation(self, value: Optional[bool] = None) -> bool:
+        ...
 
 
 @deprecated("Use the session pattern instead")

graphistry/plugins/kustograph.py

@@ -9,8 +9,8 @@
 
 from graphistry.Plottable import Plottable
 from graphistry.util import setup_logger
-from graphistry.pygraphistry import PyGraphistry
 from graphistry.plugins_types.kusto_types import KustoConfig, KustoConnectionError, KustoQueryResult, KustoSession
+from graphistry.client_session import ClientSession
 
 logger = setup_logger(__name__)
 
@@ -27,14 +27,13 @@ class KustoGraph:
     def __init__(
         self,
         *args,
-        kusto_session: Optional[KustoSession] = None,
         **kwargs: Any,
     ) -> None:
         # NOTE: Cooperative Mixin initialization passes args and kwargs along
-        kwargs['kusto_session'] = kusto_session
         super().__init__(*args, **kwargs)
 
-        self._kusto_session = kusto_session or KustoSession()
+        session = kwargs.get('pygraphistry_session', {}).get('_session', None)
+        self._kusto_session = session.kusto if isinstance(session, ClientSession) else KustoSession()
 
     def from_kusto_client(self, client: KustoClient, database: str) -> 'KustoGraph':
         self._kusto_session.client = client
@@ -173,7 +172,8 @@ def kql_graph(self, graph_name: str, snap_name: Optional[str] = None) -> Plottab
             ::
                 g = graphistry.kusto_query_graph("HoneypotNetwork").plot()
         """
-        g = self if isinstance(self, Plottable) else PyGraphistry.bind()
+        from ..plotter import Plotter
+        g = self if isinstance(self, Plottable) else Plotter()  # type: ignore
 
         if snap_name:
             graph_query = f'graph("{graph_name}", "{snap_name}" | graph-to-table nodes as N with_node_id=NodeId, edges as E with_source_id=src with_target_id=dst; N;E'
@@ -184,7 +184,7 @@ def kql_graph(self, graph_name: str, snap_name: Optional[str] = None) -> Plottab
             raise ValueError(f"Expected 2 results, got {len(results)}")
         nodes = pd.DataFrame(results[0].data, columns=results[0].column_names)
         edges = pd.DataFrame(results[1].data, columns=results[1].column_names)
-        return g.nodes(nodes, node='NodeId').edges(edges, source='src', destination='dst')
+        return g.nodes(nodes, node='NodeId').edges(edges, source='src', destination='dst')  # type: ignore
 
 
     def _kql(self, query: str) -> List[KustoQueryResult]:

graphistry/plugins/spannergraph.py

@@ -4,14 +4,14 @@
 
 import pandas as pd
 from graphistry.Plottable import Plottable
-from graphistry.pygraphistry import PyGraphistry
 from graphistry.util import setup_logger
 from graphistry.plugins_types.spanner_types import (
     SpannerConfig,
     SpannerConnectionError,
     SpannerQueryResult,
     SpannerSession,
 )
+from graphistry.client_session import ClientSession
 
 if TYPE_CHECKING:
     from google.cloud.spanner_dbapi.connection import Connection
@@ -32,13 +32,13 @@ class SpannerGraph:
     def __init__(
         self,
         *args,
-        spanner_session: Optional[SpannerSession] = None,
         **kwargs: Any,
     ) -> None:
         # NOTE: Cooperative Mixin initialization passes args and kwargs along
-        kwargs["spanner_session"] = spanner_session
         super().__init__(*args, **kwargs)
-        self._spanner_session = spanner_session or SpannerSession()
+
+        session = kwargs.get('pygraphistry_session', {}).get('_session', None)
+        self._spanner_session = session.spanner if isinstance(session, ClientSession) else SpannerSession()
 
     def from_spanner_client(self, client: Connection, database: str) -> "SpannerGraph":
         self._spanner_session.client = client
@@ -237,7 +237,8 @@ def gql_to_graph(self, query: str) -> Plottable:
                     g.plot()
      
         """
-        g = self if isinstance(self, Plottable) else PyGraphistry.bind()
+        from ..plotter import Plotter
+        g = self if isinstance(self, Plottable) else Plotter()  # type: ignore
         query_result = self._gql(query)
 
         # convert json result set to a list 
@@ -249,7 +250,7 @@ def gql_to_graph(self, query: str) -> Plottable:
         edges_df = self.get_edges_df(json_data)
 
         # TODO(tcook): add more error handling here if nodes or edges are empty
-        return g.nodes(nodes_df, 'identifier').edges(edges_df, 'source', 'destination')
+        return g.nodes(nodes_df, 'identifier').edges(edges_df, 'source', 'destination')  # type: ignore
 
     def gql_to_df(self, query: str) -> pd.DataFrame:
         """