From 61b7e61f8cfd938a3f0c1d9650e176daee637ec9 Mon Sep 17 00:00:00 2001
From: 12f23eddde <12f23eddde@gmail.com>
Date: Thu, 27 Jun 2024 21:05:41 +0800
Subject: [PATCH] feat: exclude large maps

---
 tests/test_local.py |  9 +++++++++
 woc/local.pyi       |  2 ++
 woc/local.pyx       | 10 +++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/tests/test_local.py b/tests/test_local.py
index e1fb5fa..eb02516 100644
--- a/tests/test_local.py
+++ b/tests/test_local.py
@@ -1,3 +1,4 @@
+import json
 import os
 
 import pytest
@@ -139,3 +140,11 @@ def test_version(woc):
     assert woc_u.maps == {"b2fa", "c2p", "c2dat", "b2tac"}
     woc_r = WocMapsLocal(_test_pr, version=["R"])
     assert len(woc_u.maps) + len(woc_r.maps) == len(woc.maps)
+
+
+def test_exclude_larges(woc):
+    _test_pr = os.path.join(os.path.dirname(__file__), "test_profile.json")
+    woc_nolarge = WocMapsLocal(_test_pr, exclude_larges=True)
+    assert "larges" not in json.dumps(woc_nolarge.config["maps"]), json.dumps(
+        woc_nolarge.config["maps"]
+    )
diff --git a/woc/local.pyi b/woc/local.pyi
index bb71c96..aaa5b1b 100644
--- a/woc/local.pyi
+++ b/woc/local.pyi
@@ -7,6 +7,7 @@ class WocMapsLocal(WocMapsBase):
         self,
         profile_path: Union[str, Iterable[str], None] = None,
         version: Union[str, Iterable[str], None] = None,
+        exclude_larges: bool = False,
     ) -> None:
         """
         Initialize local WoC maps with a profile.
@@ -15,6 +16,7 @@ class WocMapsLocal(WocMapsBase):
                              if not provided, use `./wocprofile.json`, `~/.wocprofile.json`, `/etc/wocprofile.json`.
         :param version: version of the profile, default to the latest version.
                         can be a single version like 'R' or a list of versions like ['R', 'U'].
+        :param exclude_larges: exclude large maps from the profile. This is useful when you don't care about files or commits appear everywhere.
         """
         ...
 
diff --git a/woc/local.pyx b/woc/local.pyx
index 1281f4e..116262d 100644
--- a/woc/local.pyx
+++ b/woc/local.pyx
@@ -535,7 +535,8 @@ def read_large(path: str, dtype: str) -> bytes:
 class WocMapsLocal(WocMapsBase):
     def __init__(self,
             profile_path: Union[str, Iterable[str], None] = None,
-            version: Union[str, Iterable[str], None] = None
+            version: Union[str, Iterable[str], None] = None,
+            exclude_larges: bool = False
         ) -> None:
         # init logger
         self._logger = logging.getLogger(__name__)
@@ -584,6 +585,13 @@ class WocMapsLocal(WocMapsBase):
             for _k in _keys_to_drop:
                 del self.config["maps"][_k]
 
+        # exclude larges
+        if exclude_larges:
+            for _k in self.config["maps"].keys():
+                for _item in self.config["maps"][_k]:
+                    if "larges" in _item:
+                        del _item["larges"]
+
         # store name of maps and objects
         self.maps = set(self.config["maps"].keys())
         self.objects = set(self.config["objects"].keys())