chore: fix formatting and types

Author: tswast

.gitignore

@@ -64,3 +64,4 @@ tests/js/node_modules/
 pylintrc
 pylintrc.test
 dummy.pkl
+.mypy_cache/

bigframes/bigquery/_operations/obj.py

@@ -13,9 +13,25 @@
 # limitations under the License.
 
 
-"""
-ObjectRef functions defined from
-https://cloud.google.com/bigquery/docs/reference/standard-sql/object-ref-functions
+"""This module integrates BigQuery built-in 'ObjectRef' functions for use with Series/DataFrame objects,
+such as OBJ.FETCH_METADATA:
+https://docs.cloud.google.com/bigquery/docs/reference/standard-sql/objectref_functions
+
+
+.. warning::
+
+    This product or feature is subject to the "Pre-GA Offerings Terms" in the
+    General Service Terms section of the `Service Specific Terms
+    <https://cloud.google.com/terms/service-terms>`_. Pre-GA products and
+    features are available "as is" and might have limited support. For more
+    information, see the `launch stage descriptions
+    <https://cloud.google.com/products?hl=en#product-launch-stages>`_.
+
+
+.. note::
+
+    To provide feedback or request support for this feature, send an email to
+    bq-objectref-feedback@google.com.
 """
 
 

bigframes/bigquery/obj.py

@@ -16,11 +16,7 @@
 <https://cloud.google.com/bigquery/docs/object-table-object-ref-functions>`_ functions.
 """
 
-from bigframes.bigquery._operations.obj import (
-    fetch_metadata,
-    get_access_url,
-    make_ref,
-)
+from bigframes.bigquery._operations.obj import fetch_metadata, get_access_url, make_ref
 
 __all__ = [
     "fetch_metadata",

bigframes/core/compile/ibis_compiler/scalar_op_registry.py

@@ -16,6 +16,7 @@
 
 import functools
 import typing
+from typing import cast
 
 from bigframes_vendored import ibis
 import bigframes_vendored.ibis.expr.api as ibis_api
@@ -1248,7 +1249,9 @@ def obj_fetch_metadata_op_impl(obj_ref: ibis_types.Value):
 @scalar_op_compiler.register_unary_op(ops.ObjGetAccessUrl, pass_op=True)
 def obj_get_access_url_op_impl(obj_ref: ibis_types.Value, op: ops.ObjGetAccessUrl):
     if op.duration is not None:
-        duration_value = ibis_types.literal(op.duration).to_interval("us")
+        duration_value = cast(
+            ibis_types.IntegerValue, ibis_types.literal(op.duration)
+        ).to_interval("us")
         return obj_get_access_url_with_duration(
             obj_ref=obj_ref, mode=op.mode, duration=duration_value
         )
@@ -2162,9 +2165,7 @@ def obj_get_access_url(obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String) -
 
 
 @ibis_udf.scalar.builtin(name="OBJ.GET_ACCESS_URL")
-def obj_get_access_url_with_duration(
-    obj_ref: _OBJ_REF_IBIS_DTYPE, mode: ibis_dtypes.String, duration: ibis_dtypes.Interval(unit="us")  # type: ignore
-) -> ibis_dtypes.JSON:  # type: ignore
+def obj_get_access_url_with_duration(obj_ref, mode, duration) -> ibis_dtypes.JSON:  # type: ignore
     """Get access url (as ObjectRefRumtime JSON) from ObjectRef."""
 
 

docs/reference/index.rst

@@ -11,6 +11,7 @@ packages.
     bigframes.bigquery
     bigframes.bigquery.ai
     bigframes.bigquery.ml
+    bigframes.bigquery.obj
     bigframes.enums
     bigframes.exceptions
     bigframes.geopandas

tests/unit/bigquery/test_obj.py

@@ -15,8 +15,6 @@
 import datetime
 from unittest.mock import MagicMock
 
-import pytest
-
 import bigframes.bigquery.obj as obj
 import bigframes.operations as ops
 import bigframes.series as series
@@ -26,26 +24,31 @@ def test_fetch_metadata_op_structure():
     op = ops.obj_fetch_metadata_op
     assert op.name == "obj_fetch_metadata"
 
+
 def test_get_access_url_op_structure():
     op = ops.ObjGetAccessUrl(mode="r")
     assert op.name == "obj_get_access_url"
     assert op.mode == "r"
     assert op.duration is None
 
+
 def test_get_access_url_with_duration_op_structure():
     op = ops.ObjGetAccessUrl(mode="rw", duration=3600000000)
     assert op.name == "obj_get_access_url"
     assert op.mode == "rw"
     assert op.duration == 3600000000
 
+
 def test_make_ref_op_structure():
     op = ops.obj_make_ref_op
     assert op.name == "obj_make_ref"
 
+
 def test_make_ref_json_op_structure():
     op = ops.obj_make_ref_json_op
     assert op.name == "obj_make_ref_json"
 
+
 def test_fetch_metadata_calls_apply_unary_op():
     s = MagicMock(spec=series.Series)
 
@@ -55,6 +58,7 @@ def test_fetch_metadata_calls_apply_unary_op():
     args, _ = s._apply_unary_op.call_args
     assert args[0] == ops.obj_fetch_metadata_op
 
+
 def test_get_access_url_calls_apply_unary_op_without_duration():
     s = MagicMock(spec=series.Series)
 
@@ -66,6 +70,7 @@ def test_get_access_url_calls_apply_unary_op_without_duration():
     assert args[0].mode == "r"
     assert args[0].duration is None
 
+
 def test_get_access_url_calls_apply_unary_op_with_duration():
     s = MagicMock(spec=series.Series)
     duration = datetime.timedelta(hours=1)
@@ -79,6 +84,7 @@ def test_get_access_url_calls_apply_unary_op_with_duration():
     # 1 hour = 3600 seconds = 3600 * 1000 * 1000 microseconds
     assert args[0].duration == 3600000000
 
+
 def test_make_ref_calls_apply_binary_op_with_authorizer():
     uri = MagicMock(spec=series.Series)
     auth = MagicMock(spec=series.Series)
@@ -90,6 +96,7 @@ def test_make_ref_calls_apply_binary_op_with_authorizer():
     assert args[0] == auth
     assert args[1] == ops.obj_make_ref_op
 
+
 def test_make_ref_calls_apply_unary_op_without_authorizer():
     json_val = MagicMock(spec=series.Series)