| # Copyright (C) 2022 The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # 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. |
| |
| import dataclasses as dc |
| from typing import Dict, List, Type, Union |
| |
| from perfetto.trace_uri_resolver import util |
| from perfetto.trace_uri_resolver.resolver import TraceContent |
| from perfetto.trace_uri_resolver.resolver import TraceGenerator |
| from perfetto.trace_uri_resolver.resolver import TraceUri |
| from perfetto.trace_uri_resolver.resolver import TraceUriResolver |
| |
| TraceReference = Union[TraceUriResolver, TraceUri, TraceContent] |
| TraceListReference = Union[TraceReference, List[TraceReference]] |
| |
| |
| class ResolverRegistry: |
| |
| @dc.dataclass |
| class Result: |
| generator: TraceGenerator |
| metadata: Dict[str, str] |
| |
| def __init__(self, resolvers: List[Type[TraceUriResolver]] = []): |
| self.registry: Dict[str, Type[TraceUriResolver]] = dict() |
| |
| for resolver in resolvers: |
| self.register(resolver) |
| |
| def register(self, provider: Type[TraceUriResolver]): |
| self.registry[provider.PREFIX] = provider |
| |
| def resolve(self, ref: TraceListReference) -> List['ResolverRegistry.Result']: |
| if isinstance(ref, list): |
| return [inner for outer in ref for inner in self._resolve_ref(outer)] |
| return self._resolve_ref(ref) |
| |
| def _resolve_ref(self, |
| ref: TraceReference) -> List['ResolverRegistry.Result']: |
| if isinstance(ref, TraceUriResolver): |
| return [ |
| _merge_metadata(outer, inner) |
| for outer in ref.resolve() |
| for inner in self.resolve(outer.trace) |
| ] |
| |
| if isinstance(ref, TraceUri): |
| return [ |
| _merge_metadata(outer, inner) |
| for outer in self._resolver_from_uri(ref).resolve() |
| for inner in self.resolve(outer.trace) |
| ] |
| |
| if hasattr(ref, 'read'): |
| return [ResolverRegistry.Result(util.read_generator(ref), {})] |
| |
| return [ResolverRegistry.Result(ref, {})] |
| |
| def _resolver_from_uri(self, uri: TraceUri) -> TraceUriResolver: |
| resolver_name, _ = util.parse_trace_uri(uri) |
| resolver_cls = self.registry[resolver_name] |
| return resolver_cls.from_trace_uri(uri) |
| |
| |
| def _merge_metadata(resolver_res: TraceUriResolver.Result, |
| registry_res: ResolverRegistry.Result |
| ) -> List[ResolverRegistry.Result]: |
| return ResolverRegistry.Result( |
| generator=registry_res.generator, |
| metadata=util.merge_dicts(resolver_res.metadata, registry_res.metadata)) |