fix all lint!

Author: Linchin

bigquery_storage/pyarrow/append_rows_with_arrow.py

@@ -13,8 +13,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from concurrent.futures import Future
 import datetime
 import decimal
+from typing import Iterable
 
 from google.cloud import bigquery
 from google.cloud import bigquery_storage_v1
@@ -158,7 +160,9 @@ def generate_pyarrow_table(num_rows: int = TABLE_LENGTH) -> pa.Table:
     return table
 
 
-def generate_write_requests(pyarrow_table):
+def generate_write_requests(
+    pyarrow_table: pa.Table,
+) -> Iterable[gapic_types.AppendRowsRequest]:
     # Determine max_chunksize of the record batches. Because max size of
     # AppendRowsRequest is 10 MB, we need to split the table if it's too big.
     # See: https://cloud.google.com/bigquery/docs/reference/storage/rpc/google.cloud.bigquery.storage.v1#appendrowsrequest
@@ -173,7 +177,9 @@ def generate_write_requests(pyarrow_table):
         yield request
 
 
-def verify_result(client, table, futures):
+def verify_result(
+    client: bigquery.Client, table: bigquery.Table, futures: "list[Future]"
+) -> None:
     bq_table = client.get_table(table)
 
     # Verify table schema.
@@ -190,7 +196,7 @@ def verify_result(client, table, futures):
     assert len(futures) == 2
 
 
-def main(project_id, dataset):
+def main(project_id: str, dataset: bigquery.Dataset) -> None:
     # Initialize clients.
     write_client = bqstorage_write_client()
     bq_client = bigquery.Client()

bigquery_storage/pyarrow/append_rows_with_arrow_test.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from google.cloud import bigquery
+
 from . import append_rows_with_arrow
 
 
-def test_append_rows_with_arrow(project_id, dataset):
+def test_append_rows_with_arrow(project_id: str, dataset: bigquery.Dataset) -> None:
     append_rows_with_arrow.main(project_id, dataset)

bigquery_storage/quickstart/quickstart.py

@@ -15,7 +15,7 @@
 import argparse
 
 
-def main(project_id="your-project-id", snapshot_millis=0):
+def main(project_id: str = "your-project-id", snapshot_millis: int = 0) -> None:
     # [START bigquerystorage_quickstart]
     from google.cloud.bigquery_storage import BigQueryReadClient, types
 

bigquery_storage/quickstart/quickstart_test.py

@@ -14,23 +14,27 @@
 
 import datetime
 
+import pytest
+
 from . import quickstart
 
 
-def now_millis():
+def now_millis() -> int:
     return int(
         (datetime.datetime.utcnow() - datetime.datetime(1970, 1, 1)).total_seconds()
         * 1000
     )
 
 
-def test_quickstart_wo_snapshot(capsys, project_id):
+def test_quickstart_wo_snapshot(capsys: pytest.CaptureFixture, project_id: str) -> None:
     quickstart.main(project_id)
     out, _ = capsys.readouterr()
     assert "unique names in states: WA" in out
 
 
-def test_quickstart_with_snapshot(capsys, project_id):
+def test_quickstart_with_snapshot(
+    capsys: pytest.CaptureFixture, project_id: str
+) -> None:
     quickstart.main(project_id, now_millis() - 5000)
     out, _ = capsys.readouterr()
     assert "unique names in states: WA" in out

bigquery_storage/snippets/append_rows_pending.py

@@ -18,10 +18,9 @@
 using the low-level generated client for Python.
 """
 
-from google.protobuf import descriptor_pb2
-
 from google.cloud import bigquery_storage_v1
 from google.cloud.bigquery_storage_v1 import types, writer
+from google.protobuf import descriptor_pb2
 
 # If you update the customer_record.proto protocol buffer definition, run:
 #
@@ -31,14 +30,14 @@
 from . import customer_record_pb2
 
 
-def create_row_data(row_num: int, name: str):
+def create_row_data(row_num: int, name: str) -> bytes:
     row = customer_record_pb2.CustomerRecord()
     row.row_num = row_num
     row.customer_name = name
     return row.SerializeToString()
 
 
-def append_rows_pending(project_id: str, dataset_id: str, table_id: str):
+def append_rows_pending(project_id: str, dataset_id: str, table_id: str) -> None:
     """Create a write stream, write some sample data, and commit the stream."""
     write_client = bigquery_storage_v1.BigQueryWriteClient()
     parent = write_client.table_path(project_id, dataset_id, table_id)

bigquery_storage/snippets/append_rows_pending_test.py

@@ -15,9 +15,8 @@
 import pathlib
 import random
 
-import pytest
-
 from google.cloud import bigquery
+import pytest
 
 from . import append_rows_pending
 
@@ -51,7 +50,7 @@ def test_append_rows_pending(
     capsys: pytest.CaptureFixture,
     bigquery_client: bigquery.Client,
     sample_data_table: str,
-):
+) -> None:
     project_id, dataset_id, table_id = sample_data_table.split(".")
     append_rows_pending.append_rows_pending(
         project_id=project_id, dataset_id=dataset_id, table_id=table_id

bigquery_storage/snippets/append_rows_proto2.py

@@ -20,10 +20,9 @@
 import datetime
 import decimal
 
-from google.protobuf import descriptor_pb2
-
 from google.cloud import bigquery_storage_v1
 from google.cloud.bigquery_storage_v1 import types, writer
+from google.protobuf import descriptor_pb2
 
 # If you make updates to the sample_data.proto protocol buffers definition,
 # run:
@@ -34,7 +33,7 @@
 from . import sample_data_pb2
 
 
-def append_rows_proto2(project_id: str, dataset_id: str, table_id: str):
+def append_rows_proto2(project_id: str, dataset_id: str, table_id: str) -> None:
     """Create a write stream, write some sample data, and commit the stream."""
     write_client = bigquery_storage_v1.BigQueryWriteClient()
     parent = write_client.table_path(project_id, dataset_id, table_id)

bigquery_storage/snippets/append_rows_proto2_test.py

@@ -17,9 +17,8 @@
 import pathlib
 import random
 
-import pytest
-
 from google.cloud import bigquery
+import pytest
 
 from . import append_rows_proto2
 
@@ -53,7 +52,7 @@ def test_append_rows_proto2(
     capsys: pytest.CaptureFixture,
     bigquery_client: bigquery.Client,
     sample_data_table: str,
-):
+) -> None:
     project_id, dataset_id, table_id = sample_data_table.split(".")
     append_rows_proto2.append_rows_proto2(
         project_id=project_id, dataset_id=dataset_id, table_id=table_id

bigquery_storage/snippets/conftest.py

@@ -12,16 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
-import test_utils.prefixer
+from typing import Generator
 
 from google.cloud import bigquery
+import pytest
+import test_utils.prefixer
 
 prefixer = test_utils.prefixer.Prefixer("python-bigquery-storage", "samples/snippets")
 
 
 @pytest.fixture(scope="session", autouse=True)
-def cleanup_datasets(bigquery_client: bigquery.Client):
+def cleanup_datasets(bigquery_client: bigquery.Client) -> None:
     for dataset in bigquery_client.list_datasets():
         if prefixer.should_cleanup(dataset.dataset_id):
             bigquery_client.delete_dataset(
@@ -30,17 +31,19 @@ def cleanup_datasets(bigquery_client: bigquery.Client):
 
 
 @pytest.fixture(scope="session")
-def bigquery_client():
+def bigquery_client() -> bigquery.Client:
     return bigquery.Client()
 
 
 @pytest.fixture(scope="session")
-def project_id(bigquery_client):
+def project_id(bigquery_client: bigquery.Client) -> str:
     return bigquery_client.project
 
 
 @pytest.fixture(scope="session")
-def dataset_id(bigquery_client: bigquery.Client, project_id: str):
+def dataset_id(
+    bigquery_client: bigquery.Client, project_id: str
+) -> Generator[str, None, None]:
     dataset_id = prefixer.create_prefix()
     full_dataset_id = f"{project_id}.{dataset_id}"
     dataset = bigquery.Dataset(full_dataset_id)
@@ -50,7 +53,9 @@ def dataset_id(bigquery_client: bigquery.Client, project_id: str):
 
 
 @pytest.fixture(scope="session")
-def dataset_id_non_us(bigquery_client: bigquery.Client, project_id: str):
+def dataset_id_non_us(
+    bigquery_client: bigquery.Client, project_id: str
+) -> Generator[str, None, None]:
     dataset_id = prefixer.create_prefix()
     full_dataset_id = f"{project_id}.{dataset_id}"
     dataset = bigquery.Dataset(full_dataset_id)

bigquery_storage/to_dataframe/read_query_results.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pandas
 
-def read_query_results():
+
+def read_query_results() -> pandas.DataFrame:
     # [START bigquerystorage_pandas_tutorial_read_query_results]
     from google.cloud import bigquery