1# Copyright (C) 2022 The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import dataclasses as dc 16from typing import Dict, List, Type, Union 17 18from perfetto.trace_uri_resolver import util 19from perfetto.trace_uri_resolver.resolver import TraceContent 20from perfetto.trace_uri_resolver.resolver import TraceGenerator 21from perfetto.trace_uri_resolver.resolver import TraceUri 22from perfetto.trace_uri_resolver.resolver import TraceUriResolver 23 24TraceReference = Union[TraceUriResolver, TraceUri, TraceContent] 25TraceListReference = Union[TraceReference, List[TraceReference]] 26 27 28class ResolverRegistry: 29 30 @dc.dataclass 31 class Result: 32 generator: TraceGenerator 33 metadata: Dict[str, str] 34 35 def __init__(self, resolvers: List[Type[TraceUriResolver]] = []): 36 self.registry: Dict[str, Type[TraceUriResolver]] = dict() 37 38 for resolver in resolvers: 39 self.register(resolver) 40 41 def register(self, provider: Type[TraceUriResolver]): 42 self.registry[provider.PREFIX] = provider 43 44 def resolve(self, ref: TraceListReference) -> List['ResolverRegistry.Result']: 45 if isinstance(ref, list): 46 return [inner for outer in ref for inner in self._resolve_ref(outer)] 47 return self._resolve_ref(ref) 48 49 def _resolve_ref(self, 50 ref: TraceReference) -> List['ResolverRegistry.Result']: 51 if isinstance(ref, TraceUriResolver): 52 return [ 53 _merge_metadata(outer, inner) 54 for outer in ref.resolve() 55 for inner in self.resolve(outer.trace) 56 ] 57 58 if isinstance(ref, TraceUri): 59 return [ 60 _merge_metadata(outer, inner) 61 for outer in self._resolver_from_uri(ref).resolve() 62 for inner in self.resolve(outer.trace) 63 ] 64 65 if hasattr(ref, 'read'): 66 return [ResolverRegistry.Result(util.read_generator(ref), {})] 67 68 return [ResolverRegistry.Result(ref, {})] 69 70 def _resolver_from_uri(self, uri: TraceUri) -> TraceUriResolver: 71 resolver_name, _ = util.parse_trace_uri(uri) 72 resolver_cls = self.registry[resolver_name] 73 return resolver_cls.from_trace_uri(uri) 74 75 76def _merge_metadata(resolver_res: TraceUriResolver.Result, 77 registry_res: ResolverRegistry.Result 78 ) -> List[ResolverRegistry.Result]: 79 return ResolverRegistry.Result( 80 generator=registry_res.generator, 81 metadata=util.merge_dicts(resolver_res.metadata, registry_res.metadata)) 82