1# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. 2# 3# This source code is licensed under the MIT license found in the 4# LICENSE file in the root directory of this source tree. 5 6import os 7import typing as tp 8import numpy as np 9from . import tools 10from . import testing 11 12 13@testing.parametrized( 14 void=([], []), 15 one=(["a"], []), 16 two=([1, 2], [(1, 2)]), 17 three=([1, 2, 3], [(1, 2), (2, 3)]), 18) 19def test_pairwise(iterator: tp.Iterable[tp.Any], expected: tp.List[tp.Tuple[tp.Any, ...]]) -> None: 20 output = list(tools.pairwise(iterator)) 21 testing.printed_assert_equal(output, expected) 22 23 24def test_roundrobin() -> None: 25 output = list(tools.roundrobin([1, 2, 3], (x for x in [4, 5, 6, 7]), (8,))) 26 np.testing.assert_array_equal(output, [1, 4, 8, 2, 5, 3, 6, 7]) 27 28 29def test_grouper() -> None: 30 output = list(tools.grouper("ABCDEFG", 3, "x")) 31 testing.printed_assert_equal(output, [list(x) for x in ["ABC", "DEF", "Gxx"]]) 32 33 34def test_sleeper() -> None: 35 min_sleep = 1e-5 36 sleeper = tools.Sleeper(min_sleep=min_sleep) 37 np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5) 38 sleeper.start_timer() 39 # not precise enough to test for exact time 40 # np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5) 41 sleeper.stop_timer() 42 # np.testing.assert_almost_equal(sleeper._get_advised_sleep_duration(), min_sleep, decimal=5) 43 44 45def test_mutable_set() -> None: 46 s: tools.OrderedSet[int] = tools.OrderedSet((1, 2, 3)) 47 assert tuple(s) == (1, 2, 3) 48 s.add(1) 49 s.add(4) 50 assert tuple(s) == (2, 3, 1, 4) 51 # 52 union = s & {3, 2} 53 assert isinstance(union, tools.OrderedSet) 54 assert tuple(union) == (2, 3) 55 # 56 assert s.pop() == 2 57 assert s.popright() == 4 58 assert tuple(s) == (3, 1) 59 # 60 s = tools.OrderedSet((1, 2, 3)) 61 intersect = s | {5, 6} 62 assert isinstance(intersect, tools.OrderedSet) 63 assert tuple(intersect) == (1, 2, 3, 5, 6) 64 intersect = {5, 6} | s 65 assert isinstance(intersect, tools.OrderedSet) 66 assert tuple(intersect) == (1, 2, 3, 5, 6) # same behavior, always appended to the end 67 68 69def test_set_env() -> None: 70 with tools.set_env(BLUBLU=1): 71 assert os.environ.get("BLUBLU", None) == "1" 72 assert os.environ.get("BLUBLU", None) is None 73 74 75def test_flatten() -> None: 76 out = tools.flatten(["a", {"truc": [4, 5]}]) 77 assert out == {"0": "a", "1.truc.0": 4, "1.truc.1": 5} 78