| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from itertools import dropwhile | |
| from time import perf_counter_ns | |
| class Timer: | |
| def __init__(self, key, d): | |
| self.key = key | |
| self.d = d | |
| if self.key not in self.d.keys(): | |
| self.d[self.key] = [] | |
| def __enter__(self): | |
| self.start = perf_counter_ns() | |
| return self | |
| def __exit__(self, *args): | |
| elapsed = perf_counter_ns() - self.start | |
| # base, unit = next( | |
| # dropwhile( | |
| # lambda t: elapsed > 10 ** (t[0] + 3), | |
| # [(0, "ns"), (3, "us"), (6, "ms"), (9, "s")], | |
| # ), | |
| # (9, "s"), | |
| # ) | |
| self.d[self.key].append(elapsed / 10**9) | |