Add all_extension_numbers_of_type in py reflection

1 files changed, 15 insertions, 10 deletions

diff --git a/src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py b/src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py
index fdb8baab0d..f9f78fcad4 100644
--- a/src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py
+++ b/src/python/grpcio_reflection/grpc_reflection/v1alpha/reflection.py
@@ -96,15 +96,20 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
         else:
             return _file_descriptor_response(descriptor)
 
-    def _extension_numbers_of_type(fully_qualified_name):
-        # TODO(atash) We're allowed to leave this unsupported according to the
-        # protocol, but we should still eventually implement it. Hits the same issue
-        # as `_file_containing_extension`, however.
-        # https://github.com/google/protobuf/issues/2248
-        return reflection_pb2.ServerReflectionResponse(
-            error_response=reflection_pb2.ErrorResponse(
-                error_code=grpc.StatusCode.UNIMPLEMENTED.value[0],
-                error_message=grpc.StatusCode.UNIMPLMENTED.value[1].encode(),))
+    def _all_extension_numbers_of_type(self, containing_type):
+        try:
+            message_descriptor = self._pool.FindMessageTypeByName(containing_type)
+            extension_numbers = tuple(sorted(
+                extension.number
+                for extension in self._pool.FindAllExtensions(message_descriptor)))
+        except KeyError:
+            return _not_found_error()
+        else:
+            return reflection_pb2.ServerReflectionResponse(
+                all_extension_numbers_response=reflection_pb2.
+                ExtensionNumberResponse(
+                    base_type_name=message_descriptor.full_name,
+                    extension_number=extension_numbers))
 
     def _list_services(self):
         return reflection_pb2.ServerReflectionResponse(
@@ -125,7 +130,7 @@ class ReflectionServicer(reflection_pb2.ServerReflectionServicer):
                     request.file_containing_extension.containing_type,
                     request.file_containing_extension.extension_number)
             elif request.HasField('all_extension_numbers_of_type'):
-                yield _all_extension_numbers_of_type(
+                yield self._all_extension_numbers_of_type(
                     request.all_extension_numbers_of_type)
             elif request.HasField('list_services'):
                 yield self._list_services()