1# BSD 3-Clause License
  2#
  3# Copyright (c) 2022-Present, nxtlo
  4# All rights reserved.
  5#
  6# Redistribution and use in source and binary forms, with or without
  7# modification, are permitted provided that the following conditions are met:
  8#
  9# * Redistributions of source code must retain the above copyright notice, this
 10#   list of conditions and the following disclaimer.
 11#
 12# * Redistributions in binary form must reproduce the above copyright notice,
 13#   this list of conditions and the following disclaimer in the documentation
 14#   and/or other materials provided with the distribution.
 15#
 16# * Neither the name of the copyright holder nor the names of its
 17#   contributors may be used to endorse or promote products derived from
 18#   this software without specific prior written permission.
 19#
 20# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 21# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 22# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 23# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 24# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 25# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 26# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 27# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 28# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 29# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 30"""An extended version of built-in `dict` but with extra functionality.
 31
 32This contains Rust's `std::collections::HashMap` methods.
 33"""
 34
 35from __future__ import annotations
 36
 37__all__ = ("HashMap", "RefMut")
 38
 39import collections.abc as collections
 40import typing
 41
 42from sain import option as option
 43from sain.convert import From
 44from sain.convert import Into
 45from sain.default import Default
 46from sain.macros import rustc_diagnostic_item
 47from sain.result import Err
 48from sain.result import Ok
 49
 50from .base_iter import Drain
 51from .base_iter import ExtractIf
 52from .base_iter import Fn
 53from .base_iter import IntoIterator
 54from .base_iter import IntoKeys
 55from .base_iter import IntoValues
 56from .base_iter import Iter
 57
 58K = typing.TypeVar("K")
 59V = typing.TypeVar("V")
 60
 61if typing.TYPE_CHECKING:
 62    from typing_extensions import Never
 63    from typing_extensions import Self
 64
 65    from sain.option import Option
 66    from sain.result import Result
 67
 68
 69class _RawMap(
 70    typing.Generic[K, V],
 71    collections.Mapping[K, V],
 72    From[collections.Iterable[tuple[K, V]]],
 73    Into[list[tuple[K, V]]],
 74):
 75    __slots__ = ("_source",)
 76
 77    def __init__(self, source: dict[K, V] | None = None, /) -> None:
 78        self._source = source if source else {}
 79
 80    # constructors
 81
 82    @classmethod
 83    def from_keys(cls, *iterable: tuple[K, V]) -> HashMap[K, V]:
 84        """Create a new `HashMap` from a sequence of key-value pairs.
 85
 86        Example
 87        -------
 88        ```py
 89        tier_list = [
 90            ("S", "Top tier"),
 91            ("A", "Very Good"),
 92            ("B", "Good"),
 93            ("C", "Average"),
 94            ("D", "dodo")
 95        ]
 96        mapping = HashMap[str, str].from_keys(*tier_list)
 97        ```
 98        """
 99        return HashMap({k: v for k, v in iterable})
100
101    # trait impls
102
103    @classmethod
104    def from_value(cls, value: collections.Iterable[tuple[K, V]]) -> HashMap[K, V]:
105        """Creates a `HashMap` from an iterable of `(K, V)` key-value paris.
106
107        Same as `HashMap.from_keys(*iter)`
108        """
109        return HashMap({k: v for k, v in value})
110
111    def into(self) -> list[tuple[K, V]]:
112        """Turn this `HashMap` into a `[(K, V); len(self)]` key-value paris.
113
114        `self` is consumed afterwards.
115
116        Example
117        -------
118        ```py
119        users = [(0, "user-0"), (1, "user-1")]
120        map1 = HashMap.from_keys(*users)
121        assert map1.into() == users
122        # map1 is dropped.
123        ```
124        """
125        return [(k, v) for k, v in self.leak().items()]
126
127    # default impls
128
129    def into_keys(self) -> IntoKeys[K]:
130        """Creates an iterator that consumes the map, yielding all of its keys.
131
132        `self` cannot be used after calling this.
133
134        Example
135        -------
136        ```py
137        map1 = HashMap({
138            "a": 1,
139            "b": 2,
140            "c": 3,
141        })
142        keys = map1.into_keys().collect()
143        assert keys == ["a", "b", "c"]
144        ```
145        """
146        return IntoKeys(self.leak().keys())
147
148    def into_values(self) -> IntoValues[V]:
149        """Creates an iterator that consumes the map, yielding all of its values.
150
151        `self` cannot be used after calling this.
152
153        Example
154        -------
155        ```py
156        map1 = HashMap({
157            "a": 1,
158            "b": 2,
159            "c": 3,
160        })
161        values = map1.into_values().collect()
162        assert values == [1, 2, 3]
163        ```
164        """
165        return IntoValues(self.leak().values())
166
167    def contains_key(self, k: K) -> bool:
168        """Check whether `k` is in `self` or not.
169
170        Example
171        -------
172        ```py
173        users = HashMap({0: "user-0"})
174        assert users.contains_key(0) is True
175        # similar to just doing
176        assert 0 in users
177        ```
178        """
179        return k in self
180
181    def is_empty(self) -> bool:
182        """Whether this `self` contains any items or not.
183
184        Example
185        -------
186        ```py
187        storage: HashMap[str, None] = HashMap()
188        if storage.is_empty():
189            ...
190        # Or just
191        if not storage:
192            ...
193        ```
194        """
195        return not self
196
197    def get(self, key: K) -> Option[V]:
198        """Get the value associated with `key`, returns `None` if not found.
199
200        Unlike `dict.get` which returns builtin `None`, This returns `Option[T]`.
201
202        Example
203        -------
204        ```py
205        users: HashMap[str, int] = HashMap()
206        assert users.get("jack").is_none()
207
208        users = HashMap({"jack": None})
209        assert users.get("jack").is_some()
210        ```
211        """
212        if key not in self:
213            return option.NOTHING
214
215        return option.Some(self._source[key])
216
217    def get_with(self, f: collections.Callable[[], K]) -> Option[V]:
218        """Get the value associated with `key` returned from a callable `f()`, returns `None` if not found.
219
220        Example
221        -------
222        ```py
223        def get_id() -> int:
224            return 0
225
226        map = HashMap({0: "buh", 1: "guh", 2: "luh"})
227        assert map.get_with(get_id).unwrap() == "a"
228        ```
229        """
230        key = f()
231        if key not in self:
232            return option.NOTHING
233
234        return option.Some(self._source[key])
235
236    def get_pairs(self, key: K) -> Option[tuple[K, V]]:
237        """Get both `key-value` pairs associated with `key`, returns `None` if not found.
238
239        Example
240        -------
241        ```py
242        users: HashMap[str, int] = HashMap()
243        assert users.get_pairs("jack").is_none()
244
245        users = HashMap({"jack": 0})
246        assert users.get("jack").unwrap() == ("jack", 0)
247        ```
248        """
249        if key not in self:
250            return option.NOTHING
251
252        return option.Some((key, self._source[key]))
253
254    def get_many(self, *keys: K) -> Option[collections.Collection[V]]:
255        """Attempts to get `len(keys)` values in the map at once.
256
257        Returns a collection of length `keys` with the results of each query.
258        None will be returned if any of the keys missing.
259
260        The time complexity is `O(N)`, where `N` is the length of `keys`.
261
262        Example
263        -------
264        ```py
265        urls = HashMap({
266            "google": "www.google.com",
267            "github": "www.github.com",
268            "facebook": "www.facebook.com",
269            "twitter": "www.twitter.com",
270        })
271        assert urls.get_many("google","github").unwrap() == ["www.google.com", "www.github.com"]
272
273        # missing keys results in `None`
274        assert urls.get_many("google", "linkedin").is_none()
275        # duplicate keys results in `None`
276        assert urls.get_many("google", "google").is_none()
277        ```
278        """
279        if not self._source:
280            return option.NOTHING
281
282        seen: set[K] = set()
283        result: list[V] = []
284
285        for key in keys:
286            if key not in self._source or key in seen:
287                return option.NOTHING
288
289            seen.add(key)
290            result.append(self._source[key])
291
292        return option.Some(result)
293
294    def iter(self) -> Iter[tuple[K, V]]:
295        """Creates an iterator over the key-value pairs of the map.
296
297        Example
298        -------
299        ```py
300        map = HashMap.from_keys([
301            ("a", 1),
302            ("b", 2),
303        ])
304        for k, v in map.iter():
305            print(f"{k=}: {v=}")
306        ```
307        """
308        return Iter(self._source.items())
309
310    def into_iter(self) -> IntoIterator[K, V]:
311        """Creates a consuming iterator, that is, one that moves each key-value pair out of the map.
312        The map cannot be used after calling this.
313
314        Example
315        -------
316        ```py
317        map = HashMap.from_keys([
318            ("a", 1),
319            ("b", 2),
320            ("c", 3),
321        ])
322        for k, v in map.into_iter():
323            print(f"{k=}: {v=}")
324        ```
325        """
326        return IntoIterator(self.leak())
327
328    def len(self) -> int:
329        """Get how many elements are in this map.
330
331        Example
332        -------
333        ```py
334        map: HashMap[str, int] = HashMap()
335        assert map.len() == 0
336        ```
337        """
338        return len(self._source)
339
340    # conversions
341
342    def leak(self) -> collections.MutableMapping[K, V]:
343        """Leaks and returns a mutable reference to the underlying map.
344
345        `self` becomes unusable after calling this.
346
347        Example
348        -------
349        ```py
350        map = HashMap({0: 1})
351        inner = map.leak()
352        assert inner == {0: 1}
353        ```
354        """
355        out = self._source
356        del self._source
357        return out
358
359    # built-ins
360
361    def copy(self) -> Self:
362        """Copy the contents of this hash map into a new one.
363
364        Example
365        -------
366        ```py
367        hashmap: HashMap[str, None] = HashMap({'1': None, '2': None})
368        copy = hashmap.copy()
369        assert hashmap == copy
370        ```
371        """
372        return self.__class__(self._source.copy())
373
374    def __repr__(self) -> str:
375        return self._source.__repr__()
376
377    def __iter__(self) -> collections.Iterator[K]:
378        return self._source.__iter__()
379
380    def __len__(self) -> int:
381        return self._source.__len__()
382
383    def __getitem__(self, key: K, /) -> V:
384        return self._source[key]
385
386    def __setitem__(self, _key: Never, _value: Never) -> Never:
387        raise NotImplementedError(
388            "`HashMap` is immutable, use `.as_mut()` to make it mutable instead"
389        )
390
391    def __delitem__(self, _key: Never) -> Never:
392        raise NotImplementedError(
393            "`HashMap` is immutable, use `.as_mut()` to make it mutable instead"
394        )
395
396
397@rustc_diagnostic_item("HashMap")
398class HashMap(_RawMap[K, V], Default["HashMap[K, V]"]):
399    """An immutable key-value dictionary.
400
401    But default, `HashMap` is immutable, it cannot change its values after initializing it. however,
402    you can return a mutable reference to this hashmap via `HashMap.as_mut` method, it returns a reference to the underlying map.
403
404    Example
405    -------
406    ```py
407    # initialize it with a source. after that, item insertion is not allowed.
408    books: HashMap[str, str] = HashMap({
409        "Adventures of Huckleberry Finn": "My favorite book."
410    })
411    # get a mutable reference to it to be able to.
412    books_mut = books.as_mut()
413    # You can either call `.insert`, which is similar to Rust's.
414    books_mut.insert(
415        "Grimms Fairy Tales",
416        "Masterpiece.",
417    )
418    # Or via item assignments
419    books_mut["Pride and Prejudice"] = "Very enjoyable."
420    print(books)
421
422    for book, review in books.items():
423        print(book, review)
424    ```
425
426    Parameters
427    ----------
428    source: `dict[K, V]`
429        A dictionary to point to. this will be used as the underlying source.
430    """
431
432    __slots__ = ("_source",)
433
434    def __init__(self, source: dict[K, V] | None = None, /) -> None:
435        super().__init__(source)
436
437    @staticmethod
438    def default() -> HashMap[K, V]:
439        """Creates an empty `HashMap<K, V>`."""
440        return HashMap()
441
442    @classmethod
443    def from_mut(cls, source: dict[K, V] | None = None, /) -> RefMut[K, V]:
444        """Create a new mutable `HashMap`, with the given source if available.
445
446        Example
447        -------
448        ```py
449        books = HashMap.from_mut()
450        books[0] = "Twilight"
451        ```
452        """
453        return RefMut(source or {})
454
455    @staticmethod
456    def from_keys_mut(*iterable: tuple[K, V]) -> RefMut[K, V]:
457        """Create a new mutable `HashMap` from a sequence of key-value pairs.
458
459        Example
460        -------
461        ```py
462        tier_list = [
463            ("S", "Top tier"),
464            ("A", "Very Good"),
465            ("B", "Good"),
466            ("C", "Average"),
467            ("D", "dodo")
468        ]
469        mapping: RefMut[str, str] = HashMap.from_keys_mut(*tier_list)
470        ```
471        """
472        return RefMut({k: v for k, v in iterable})
473
474    def as_mut(self) -> RefMut[K, V]:
475        """Get a mutable reference to this hash map.
476
477        Example
478        -------
479        ```py
480        map: HashMap[str, float] = HashMap()
481
482        # Get a reference to map
483        mut = map.as_mut()
484        mut.insert("sqrt", 1.0)
485        del mut # not needed anymore
486
487        assert map == {"sqrt": 1.0}
488        ```
489        """
490        return RefMut(self._source)
491
492
493@typing.final
494class RefMut(_RawMap[K, V], collections.MutableMapping[K, V], Default["RefMut[K, V]"]):
495    """A reference to a mutable dictionary / hashmap.
496
497    Ideally, you want to use `HashMap.as_mut()` / `HashMap.from_mut` methods to create this.
498    """
499
500    __slots__ = ("_source",)
501
502    def __init__(self, source: dict[K, V], /) -> None:
503        super().__init__(source)
504        self._source = source
505
506    @staticmethod
507    def default() -> RefMut[K, V]:
508        """Creates a new, mutable, empty `RefMut<K, V>`."""
509        return RefMut({})
510
511    def insert(self, key: K, value: V) -> Option[V]:
512        """Insert a key/value pair into the map.
513
514        if `key` is not present in the map, `None` is returned. otherwise, the value is updated, and the old value
515        is returned.
516
517        Example
518        -------
519        ```py
520        users = HashMap.from_mut({0: "admin"})
521        assert users.insert(1, "admin").is_none()
522        old = users.insert(0, "normal").unwrap()
523        assert old == "admin"
524        ```
525        """
526        if key not in self:
527            self[key] = value
528            return option.NOTHING
529
530        old = self[key]
531        self[key] = value
532        return option.Some(old)
533
534    def try_insert(self, key: K, value: V) -> Result[V, K]:
535        """Tries to insert `key` and `value` into the map, returning `Err(key)` if `key` is already present.
536
537        Example
538        -------
539        ```py
540        users = HashMap.from_mut({0: "admin"})
541        assert users.try_insert(1, "claudia").is_ok()
542        assert users.try_insert(0, "doppler").is_err()
543        ```
544        """
545        if key in self:
546            return Err(key)
547
548        self[key] = value
549        return Ok(value)
550
551    def remove(self, key: K) -> Option[V]:
552        """Remove a key from the map, returning the value of the key if it was previously in the map.
553
554        Example
555        -------
556        ```py
557        map = HashMap.from_mut()
558        map.insert(0, "a")
559        map.remove(0).unwrap() == "a"
560        map.remove(0).is_none()
561        ```
562        """
563        if key not in self:
564            return option.NOTHING
565
566        v = self[key]
567        del self[key]
568        return option.Some(v)
569
570    def remove_entry(self, key: K) -> Option[tuple[K, V]]:
571        """Remove a key from the map, returning the key and value if it was previously in the map.
572
573        Example
574        -------
575        ```py
576        map = HashMap.from_mut()
577        map[0] = "a"
578        assert map.remove_entry(0).unwrap() == (0, "a")
579        assert map.remove_entry(0).is_none()
580        ```
581        """
582        if key not in self:
583            return option.NOTHING
584
585        v = self[key]
586        del self[key]
587        return option.Some((key, v))
588
589    def retain(self, f: collections.Callable[[K, V], bool]) -> None:
590        """Remove items from this map based on `f(key, value)` returning `False`.
591
592        Example
593        -------
594        ```py
595        users = HashMap.from_mut({
596            "user1": "admin",
597            "user2": "admin",
598            "user3": "regular",
599            "jack": "admin"
600        })
601        users.retain(
602            lambda user, role: role == "admin" and
603            user.startswith("user")
604        )
605        for user, role in users.items():
606            print(user, role)
607
608        # user1 admin
609        # user2 admin
610        """
611        for k, v in self._source.copy().items():
612            if not f(k, v):
613                del self[k]
614
615    def drain(self) -> Drain[K, V]:
616        """Clears the map, returning all key-value pairs as an iterator. Keeps the hashmap for reuse.
617
618        Example
619        -------
620        ```py
621        map = HashMap.from_mut()
622        map.insert(0, "a")
623        map.drain().collect() == [(0, "a")]
624        assert map.is_empty()
625        ```
626        """
627        return Drain(self._source)
628
629    def extract_if(self, pred: Fn[K, V]) -> ExtractIf[K, V]:
630        """Creates an iterator which uses a closure to determine if an element should be removed.
631
632        If the closure returns true, the element is removed from the map and yielded. If the closure returns false,
633        the element remains in the map and will not be yielded.
634
635        If you don't need at iterator, use `retain` instead.
636
637        Example
638        -------
639        ```py
640        odds = HashMap.from_mut({0: 0, 1: 1, 2: 2, 3: 3, 4: 4})
641        evens = odds.extract_if(lambda k, _v: k % 2 == 0).collect()
642        assert odds == {1: 1, 3: 3}
643        assert evens == [(0, 0), (2, 2), (4, 4)]
644        ```
645        """
646        return ExtractIf(self._source, pred)
647
648    def __repr__(self) -> str:
649        return self._source.__repr__()
650
651    def __setitem__(self, key: K, value: V, /) -> None:
652        self._source[key] = value
653
654    def __delitem__(self, key: K, /) -> None:
655        del self._source[key]
656
657    def __iter__(self) -> collections.Iterator[K]:
658        return self._source.__iter__()
659
660    def __len__(self) -> int:
661        return self._source.__len__()