diff --git a/gallery_dl/util.py b/gallery_dl/util.py index 7d1cc32de7..657f66ba22 100644 --- a/gallery_dl/util.py +++ b/gallery_dl/util.py @@ -71,6 +71,24 @@ def unique_sequence(iterable): yield element +def contains(values, elements): + """Returns True if at least one of 'elements' is contained in 'values'""" + if not isinstance(elements, (tuple, list)): + elements = (elements,) + + if isinstance(values, str): + fmt = r"\b{}\b".format + for e in elements: + if re.search(fmt(e), values): + return True + else: + for e in elements: + if e in values: + return True + + return False + + def raises(cls): """Returns a function that raises 'cls' as exception""" def wrap(*args): @@ -429,6 +447,7 @@ def __str__(): SENTINEL = object() SPECIAL_EXTRACTORS = {"oauth", "recursive", "test"} GLOBALS = { + "contains" : contains, "parse_int": text.parse_int, "urlsplit" : urllib.parse.urlsplit, "datetime" : datetime.datetime, diff --git a/test/test_util.py b/test/test_util.py index d5b9784f92..b7684601dc 100644 --- a/test/test_util.py +++ b/test/test_util.py @@ -319,6 +319,23 @@ def test_unique_sequence(self): self.assertSequenceEqual( list(util.unique_sequence([1, 2, 1, 3, 2, 1])), [1, 2, 1, 3, 2, 1]) + def test_contains(self): + c = [1, "2", 3, 4, "5", "foo"] + self.assertTrue(util.contains(c, 1)) + self.assertTrue(util.contains(c, "foo")) + self.assertTrue(util.contains(c, [1, 3, "5"])) + self.assertTrue(util.contains(c, ["a", "b", "5"])) + self.assertFalse(util.contains(c, "bar")) + self.assertFalse(util.contains(c, [2, 5, "bar"])) + + s = "1, 2, 3, asd, qwe, yxc, foo, bar" + self.assertTrue(util.contains(s, 2)) + self.assertTrue(util.contains(s, "asd")) + self.assertTrue(util.contains(s, ["asd", "qwe", "yxc"])) + self.assertTrue(util.contains(s, ["sdf", "dfg", "qwe"])) + self.assertFalse(util.contains(s, "tag1")) + self.assertFalse(util.contains(s, ["tag1", "tag2", "tag3"])) + def test_raises(self): func = util.raises(Exception) with self.assertRaises(Exception):