Merge pull request #2642 from dhermes/register-type-optional

Make type_url optional when registering types.

bigtable/google/cloud/bigtable/cluster.py

@@ -22,8 +22,7 @@
 from google.cloud.bigtable._generated import (
     bigtable_instance_admin_pb2 as messages_v2_pb2)
 from google.cloud.operation import Operation
-from google.cloud.operation import _compute_type_url
-from google.cloud.operation import register_type_url
+from google.cloud.operation import register_type
 
 
 _CLUSTER_NAME_RE = re.compile(r'^projects/(?P<project>[^/]+)/'
@@ -34,10 +33,7 @@
 """Default number of nodes to use when creating a cluster."""
 
 
-_UPDATE_CLUSTER_METADATA_URL = _compute_type_url(
-    messages_v2_pb2.UpdateClusterMetadata)
-register_type_url(
-    _UPDATE_CLUSTER_METADATA_URL, messages_v2_pb2.UpdateClusterMetadata)
+register_type(messages_v2_pb2.UpdateClusterMetadata)
 
 
 def _prepare_create_request(cluster):

bigtable/google/cloud/bigtable/instance.py

@@ -27,21 +27,16 @@
 from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES
 from google.cloud.bigtable.table import Table
 from google.cloud.operation import Operation
-from google.cloud.operation import _compute_type_url
-from google.cloud.operation import register_type_url
+from google.cloud.operation import register_type
 
 
 _EXISTING_INSTANCE_LOCATION_ID = 'see-existing-cluster'
 _INSTANCE_NAME_RE = re.compile(r'^projects/(?P<project>[^/]+)/'
                                r'instances/(?P<instance_id>[a-z][-a-z0-9]*)$')
 
 
-_CREATE_INSTANCE_METADATA_URL = _compute_type_url(
-    messages_v2_pb2.CreateInstanceMetadata)
-register_type_url(
-    _CREATE_INSTANCE_METADATA_URL, messages_v2_pb2.CreateInstanceMetadata)
-_INSTANCE_METADATA_URL = _compute_type_url(data_v2_pb2.Instance)
-register_type_url(_INSTANCE_METADATA_URL, data_v2_pb2.Instance)
+register_type(messages_v2_pb2.CreateInstanceMetadata)
+register_type(data_v2_pb2.Instance)
 
 
 def _prepare_create_request(instance):

bigtable/unit_tests/test_cluster.py

@@ -283,7 +283,6 @@ def test_update(self):
         from google.cloud.bigtable._generated import (
             bigtable_instance_admin_pb2 as messages_v2_pb2)
         from unit_tests._testing import _FakeStub
-        from google.cloud.bigtable.cluster import _UPDATE_CLUSTER_METADATA_URL
 
         NOW = datetime.datetime.utcnow()
         NOW_PB = _datetime_to_pb_timestamp(NOW)
@@ -307,10 +306,12 @@ def test_update(self):
             'operations/projects/%s/instances/%s/clusters/%s/operations/%d' %
             (self.PROJECT, self.INSTANCE_ID, self.CLUSTER_ID, OP_ID))
         metadata = messages_v2_pb2.UpdateClusterMetadata(request_time=NOW_PB)
+        type_url = 'type.googleapis.com/%s' % (
+            messages_v2_pb2.UpdateClusterMetadata.DESCRIPTOR.full_name,)
         response_pb = operations_pb2.Operation(
             name=OP_NAME,
             metadata=Any(
-                type_url=_UPDATE_CLUSTER_METADATA_URL,
+                type_url=type_url,
                 value=metadata.SerializeToString()
             )
         )

bigtable/unit_tests/test_instance.py

@@ -236,8 +236,6 @@ def test_create(self):
         from unit_tests._testing import _FakeStub
         from google.cloud.operation import Operation
         from google.cloud.bigtable.cluster import DEFAULT_SERVE_NODES
-        from google.cloud.bigtable.instance import (
-            _CREATE_INSTANCE_METADATA_URL)
 
         NOW = datetime.datetime.utcnow()
         NOW_PB = _datetime_to_pb_timestamp(NOW)
@@ -247,10 +245,12 @@ def test_create(self):
 
         # Create response_pb
         metadata = messages_v2_pb2.CreateInstanceMetadata(request_time=NOW_PB)
+        type_url = 'type.googleapis.com/%s' % (
+            messages_v2_pb2.CreateInstanceMetadata.DESCRIPTOR.full_name,)
         response_pb = operations_pb2.Operation(
             name=self.OP_NAME,
             metadata=Any(
-                type_url=_CREATE_INSTANCE_METADATA_URL,
+                type_url=type_url,
                 value=metadata.SerializeToString(),
                 )
             )

core/google/cloud/operation.py

@@ -39,17 +39,20 @@ def _compute_type_url(klass, prefix=_GOOGLE_APIS_PREFIX):
     return '%s/%s' % (prefix, name)
 
 
-def register_type_url(type_url, klass):
+def register_type(klass, type_url=None):
     """Register a klass as the factory for a given type URL.
 
-    :type type_url: str
-    :param type_url: URL naming the type
-
     :type klass: type
     :param klass: class to be used as a factory for the given type
 
+    :type type_url: str
+    :param type_url: (Optional) URL naming the type. If not provided,
+                     infers the URL from the type descriptor.
+
     :raises: ValueError if a registration already exists for the URL.
     """
+    if type_url is None:
+        type_url = _compute_type_url(klass)
     if type_url in _TYPE_URL_MAP:
         if _TYPE_URL_MAP[type_url] is not klass:
             raise ValueError("Conflict: %s" % (_TYPE_URL_MAP[type_url],))
@@ -123,7 +126,7 @@ class Operation(object):
     """Metadata about the current operation (as a protobuf).
 
     Code that uses operations must register the metadata types (via
-    :func:`register_type_url`) to ensure that the metadata fields can be
+    :func:`register_type`) to ensure that the metadata fields can be
     converted into the correct types.
     """
 

core/unit_tests/test_operation.py

@@ -44,48 +44,63 @@ def test_w_prefix(self):
             '%s/%s' % (PREFIX, Struct.DESCRIPTOR.full_name))
 
 
-class Test_register_type_url(unittest.TestCase):
+class Test_register_type(unittest.TestCase):
 
-    def _callFUT(self, type_url, klass):
-        from google.cloud.operation import register_type_url
-        register_type_url(type_url, klass)
+    def _callFUT(self, klass, type_url=None):
+        from google.cloud.operation import register_type
+        register_type(klass, type_url=type_url)
 
-    def test_simple(self):
+    def test_explicit(self):
         from google.cloud import operation as MUT
         from google.cloud._testing import _Monkey
-        TYPE_URI = 'testing.google-cloud-python.com/testing'
+
+        type_url = 'testing.google-cloud-python.com/testing'
         klass = object()
         type_url_map = {}
 
         with _Monkey(MUT, _TYPE_URL_MAP=type_url_map):
-            self._callFUT(TYPE_URI, klass)
+            self._callFUT(klass, type_url)
+
+        self.assertEqual(type_url_map, {type_url: klass})
 
-        self.assertEqual(type_url_map, {TYPE_URI: klass})
+    def test_default(self):
+        from google.protobuf.struct_pb2 import Struct
+        from google.cloud._testing import _Monkey
+        from google.cloud import operation as MUT
+
+        type_url_map = {}
+        with _Monkey(MUT, _TYPE_URL_MAP=type_url_map):
+            self._callFUT(Struct)
+
+        type_url = MUT._compute_type_url(Struct)
+        self.assertEqual(type_url_map, {type_url: Struct})
 
     def test_w_same_class(self):
         from google.cloud import operation as MUT
         from google.cloud._testing import _Monkey
-        TYPE_URI = 'testing.google-cloud-python.com/testing'
+
+        type_url = 'testing.google-cloud-python.com/testing'
         klass = object()
-        type_url_map = {TYPE_URI: klass}
+        type_url_map = {type_url: klass}
 
         with _Monkey(MUT, _TYPE_URL_MAP=type_url_map):
-            self._callFUT(TYPE_URI, klass)
+            self._callFUT(klass, type_url)
 
-        self.assertEqual(type_url_map, {TYPE_URI: klass})
+        self.assertEqual(type_url_map, {type_url: klass})
 
     def test_w_conflict(self):
         from google.cloud import operation as MUT
         from google.cloud._testing import _Monkey
-        TYPE_URI = 'testing.google-cloud-python.com/testing'
+
+        type_url = 'testing.google-cloud-python.com/testing'
         klass, other = object(), object()
-        type_url_map = {TYPE_URI: other}
+        type_url_map = {type_url: other}
 
         with _Monkey(MUT, _TYPE_URL_MAP=type_url_map):
             with self.assertRaises(ValueError):
-                self._callFUT(TYPE_URI, klass)
+                self._callFUT(klass, type_url)
 
-        self.assertEqual(type_url_map, {TYPE_URI: other})
+        self.assertEqual(type_url_map, {type_url: other})
 
 
 class OperationTests(unittest.TestCase):