athena-intelligence 0.1.107__tar.gz → 0.1.109__tar.gz

{athena_intelligence-0.1.107 → athena_intelligence-0.1.109}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: athena-intelligence
-Version: 0.1.107
+Version: 0.1.109
 Summary: Athena Intelligence Python Library
 Requires-Python: >=3.8,<4.0
 Classifier: Intended Audience :: Developers

{athena_intelligence-0.1.107 → athena_intelligence-0.1.109}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "athena-intelligence"
-version = "0.1.107"
+version = "0.1.109"
 description = "Athena Intelligence Python Library"
 readme = "README.md"
 authors = []

{athena_intelligence-0.1.107 → athena_intelligence-0.1.109}/src/athena/client.py

@@ -48,7 +48,7 @@ class AthenaAsset:
     data: bytes
     media_type: str
 
-    def _repr_mimebundle_(self):
+    def _repr_mimebundle_(self, include=None, exclude=None):
         if self.media_type == "application/sql":
             # it is safe to import IPython in `_repr_mimebundle_`
             # as this is only intended to be invoked by IPython.
@@ -213,9 +213,8 @@ class WrappedToolsClient(ToolsClient):
             "application/sql",
         }
 
-        data = file_bytes.read()
-
         if media_type in supported_media_types:
+            data = file_bytes.read()
             return AthenaAsset(
                 asset_id=asset_id,
                 data=data,
@@ -223,7 +222,7 @@ class WrappedToolsClient(ToolsClient):
             )
 
         if media_type in _pandas_media_types:
-            return _to_pandas_df(file_bytes)
+            return _to_pandas_df(file_bytes, media_type=media_type)
 
         raise NotImplementedError("Assets of `{media_type}` type are not yet supported")
     
@@ -488,21 +487,27 @@ _pandas_media_types = {
 }
 
 
-def _to_pandas_df(bytes_io: io.BytesIO, *args, **kwargs):
+def _to_pandas_df(
+    bytes_io: io.BytesIO, *args, media_type: Union[str, None] = None, **kwargs
+):
     import pandas as pd
 
-    media_type = _infer_media_type(bytes_io)
+    if media_type is None:
+        media_type = _infer_media_type(bytes_io)
 
-    if media_type == 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet':
-        return pd.read_excel(bytes_io, *args, engine='openpyxl', **kwargs)
-    elif media_type in {'application/vnd.apache.parquet', 'application/octet-stream'}:
+    if (
+        media_type
+        == "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+    ):
+        return pd.read_excel(bytes_io, *args, engine="openpyxl", **kwargs)
+    elif media_type in {"application/vnd.apache.parquet", "application/octet-stream"}:
         return pd.read_parquet(bytes_io, *args, **kwargs)
-    elif media_type == 'application/vnd.ms-excel':
+    elif media_type == "application/vnd.ms-excel":
         return pd.read_excel(bytes_io, *args, **kwargs)
-    elif media_type in {'text/csv', 'application/csv'}:
+    elif media_type in {"text/csv", "application/csv"}:
         return pd.read_csv(bytes_io, *args, **kwargs)
     else:
-        raise Exception("Unknown media type")
+        raise Exception(f"Unknown media type: {media_type}")
 
 
 def _convert_asset_object(

{athena_intelligence-0.1.107 → athena_intelligence-0.1.109}/src/athena/core/client_wrapper.py

@@ -17,7 +17,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.107",
+            "X-Fern-SDK-Version": "0.1.109",
         }
         headers["X-API-KEY"] = self.api_key
         return headers