diff --git a/google/cloud/bigquery/model.py b/google/cloud/bigquery/model.py index 1143b71f9..6fdfc5596 100644 --- a/google/cloud/bigquery/model.py +++ b/google/cloud/bigquery/model.py @@ -18,7 +18,7 @@ import copy -from google.protobuf import json_format +import json import six import google.cloud._helpers @@ -55,7 +55,7 @@ class Model(object): def __init__(self, model_ref): # Use _proto on read-only properties to use it's built-in type # conversion. - self._proto = types.Model()._pb + self._proto = types.Model() # Use _properties on read-write properties to match the REST API # semantics. The BigQuery API makes a distinction between an unset @@ -67,7 +67,7 @@ def __init__(self, model_ref): model_ref = ModelReference.from_string(model_ref) if model_ref: - self._proto.model_reference.CopyFrom(model_ref._proto) + self._proto.model_reference = model_ref._proto @property def reference(self): @@ -305,9 +305,7 @@ def from_api_repr(cls, resource): start_time = datetime_helpers.from_microseconds(1e3 * float(start_time)) training_run["startTime"] = datetime_helpers.to_rfc3339(start_time) - this._proto = json_format.ParseDict( - resource, types.Model()._pb, ignore_unknown_fields=True - ) + this._proto = types.Model.from_json(json.dumps(resource)) return this def _build_resource(self, filter_fields): @@ -323,7 +321,7 @@ def to_api_repr(self): Returns: Dict[str, object]: Model reference represented as an API resource """ - return json_format.MessageToDict(self._proto) + return json.loads(types.Model.to_json(self._proto)) class ModelReference(object): @@ -334,7 +332,7 @@ class ModelReference(object): """ def __init__(self): - self._proto = types.ModelReference()._pb + self._proto = types.ModelReference() self._properties = {} @property @@ -377,9 +375,7 @@ def from_api_repr(cls, resource): # Keep a reference to the resource as a workaround to find unknown # field values. ref._properties = resource - ref._proto = json_format.ParseDict( - resource, types.ModelReference()._pb, ignore_unknown_fields=True - ) + ref._proto = types.ModelReference.from_json(json.dumps(resource)) return ref @@ -418,7 +414,7 @@ def to_api_repr(self): Returns: Dict[str, object]: Model reference represented as an API resource """ - return json_format.MessageToDict(self._proto) + return json.loads(types.ModelReference.to_json(self._proto)) def _key(self): """Unique key for this model. diff --git a/tests/unit/model/test_model.py b/tests/unit/model/test_model.py index 9fa29a496..1d7605d54 100644 --- a/tests/unit/model/test_model.py +++ b/tests/unit/model/test_model.py @@ -123,21 +123,21 @@ def test_from_api_repr(target_class): assert got.training_runs[0].training_options.initial_learn_rate == 1.0 assert ( got.training_runs[0] - .start_time.ToDatetime() + .start_time .replace(tzinfo=google.cloud._helpers.UTC) == creation_time ) assert got.training_runs[1].training_options.initial_learn_rate == 0.5 assert ( got.training_runs[1] - .start_time.ToDatetime() + .start_time .replace(tzinfo=google.cloud._helpers.UTC) == modified_time ) assert got.training_runs[2].training_options.initial_learn_rate == 0.25 assert ( got.training_runs[2] - .start_time.ToDatetime() + .start_time .replace(tzinfo=google.cloud._helpers.UTC) == expiration_time ) @@ -321,7 +321,7 @@ def test_repr(target_class): def test_to_api_repr(target_class): - from google.protobuf import json_format + import json model = target_class("my-proj.my_dset.my_model") resource = { @@ -357,8 +357,6 @@ def test_to_api_repr(target_class): "kmsKeyName": "projects/1/locations/us/keyRings/1/cryptoKeys/1" }, } - model._proto = json_format.ParseDict( - resource, types.Model()._pb, ignore_unknown_fields=True - ) + model._proto = types.Model.from_json(json.dumps(resource)) got = model.to_api_repr() assert got == resource