Fix duplicates in random/wikipedia generator and adding timeout to wikipedia requests

Author: AghilesAzzoug

trdg/generators/from_random.py

@@ -10,41 +10,42 @@ class GeneratorFromRandom:
     """Generator that uses randomly generated words"""
 
     def __init__(
-        self,
-        count=-1,
-        length=1,
-        allow_variable=False,
-        use_letters=True,
-        use_numbers=True,
-        use_symbols=True,
-        fonts=[],
-        language="en",
-        size=32,
-        skewing_angle=0,
-        random_skew=False,
-        blur=0,
-        random_blur=False,
-        background_type=0,
-        distorsion_type=0,
-        distorsion_orientation=0,
-        is_handwritten=False,
-        width=-1,
-        alignment=1,
-        text_color="#282828",
-        orientation=0,
-        space_width=1.0,
-        character_spacing=0,
-        margins=(5, 5, 5, 5),
-        fit=False,
-        output_mask=False,
-        word_split=False,
-        image_dir=os.path.join(
-            "..", os.path.split(os.path.realpath(__file__))[0], "images"
-        ),
-        stroke_width=0, 
-        stroke_fill="#282828",
-        image_mode="RGB",
+            self,
+            count=-1,
+            length=1,
+            allow_variable=False,
+            use_letters=True,
+            use_numbers=True,
+            use_symbols=True,
+            fonts=[],
+            language="en",
+            size=32,
+            skewing_angle=0,
+            random_skew=False,
+            blur=0,
+            random_blur=False,
+            background_type=0,
+            distorsion_type=0,
+            distorsion_orientation=0,
+            is_handwritten=False,
+            width=-1,
+            alignment=1,
+            text_color="#282828",
+            orientation=0,
+            space_width=1.0,
+            character_spacing=0,
+            margins=(5, 5, 5, 5),
+            fit=False,
+            output_mask=False,
+            word_split=False,
+            image_dir=os.path.join(
+                "..", os.path.split(os.path.realpath(__file__))[0], "images"
+            ),
+            stroke_width=0,
+            stroke_fill="#282828",
+            image_mode="RGB",
     ):
+        self.generated_count = 0
         self.count = count
         self.length = length
         self.allow_variable = allow_variable
@@ -91,9 +92,12 @@ def __init__(
         )
 
     def __iter__(self):
-        return self.generator
+        return self
 
     def __next__(self):
+        if self.generated_count == self.count:
+            raise StopIteration
+        self.generated_count += 1
         return self.next()
 
     def next(self):

trdg/generators/from_wikipedia.py

@@ -10,37 +10,38 @@ class GeneratorFromWikipedia:
     """Generator that uses sentences taken from random Wikipedia articles"""
 
     def __init__(
-        self,
-        count=-1,
-        minimum_length=1,
-        fonts=[],
-        language="en",
-        size=32,
-        skewing_angle=0,
-        random_skew=False,
-        blur=0,
-        random_blur=False,
-        background_type=0,
-        distorsion_type=0,
-        distorsion_orientation=0,
-        is_handwritten=False,
-        width=-1,
-        alignment=1,
-        text_color="#282828",
-        orientation=0,
-        space_width=1.0,
-        character_spacing=0,
-        margins=(5, 5, 5, 5),
-        fit=False,
-        output_mask=False,
-        word_split=False,
-        image_dir=os.path.join(
-            "..", os.path.split(os.path.realpath(__file__))[0], "images"
-        ),
-        stroke_width=0, 
-        stroke_fill="#282828",
-        image_mode="RGB", 
+            self,
+            count=-1,
+            minimum_length=1,
+            fonts=[],
+            language="en",
+            size=32,
+            skewing_angle=0,
+            random_skew=False,
+            blur=0,
+            random_blur=False,
+            background_type=0,
+            distorsion_type=0,
+            distorsion_orientation=0,
+            is_handwritten=False,
+            width=-1,
+            alignment=1,
+            text_color="#282828",
+            orientation=0,
+            space_width=1.0,
+            character_spacing=0,
+            margins=(5, 5, 5, 5),
+            fit=False,
+            output_mask=False,
+            word_split=False,
+            image_dir=os.path.join(
+                "..", os.path.split(os.path.realpath(__file__))[0], "images"
+            ),
+            stroke_width=0,
+            stroke_fill="#282828",
+            image_mode="RGB",
     ):
+        self.generated_count = 0
         self.count = count
         self.minimum_length = minimum_length
         self.language = language
@@ -75,9 +76,12 @@ def __init__(
         )
 
     def __iter__(self):
-        return self.generator
+        return self
 
     def __next__(self):
+        if self.generated_count == self.count:
+            raise StopIteration
+        self.generated_count += 1
         return self.next()
 
     def next(self):

trdg/string_generator.py

@@ -19,7 +19,7 @@ def create_strings_from_file(filename, count):
             raise Exception("No lines could be read in file")
         while len(strings) < count:
             if len(lines) >= count - len(strings):
-                strings.extend(lines[0 : count - len(strings)])
+                strings.extend(lines[0: count - len(strings)])
             else:
                 strings.extend(lines)
 
@@ -50,7 +50,12 @@ def create_strings_from_wikipedia(minimum_length, count, lang):
 
     while len(sentences) < count:
         # We fetch a random page
-        page = requests.get("https://{}.wikipedia.org/wiki/Special:Random".format(lang))
+
+        page_url = "https://{}.wikipedia.org/wiki/Special:Random".format(lang)
+        try:
+            page = requests.get(page_url, timeout=3.0)  # take into account timeouts
+        except requests.exceptions.Timeout:
+            continue
 
         soup = BeautifulSoup(page.text, "html.parser")
 
@@ -61,8 +66,8 @@ def create_strings_from_wikipedia(minimum_length, count, lang):
         lines = list(
             filter(
                 lambda s: len(s.split(" ")) > minimum_length
-                and not "Wikipedia" in s
-                and not "wikipedia" in s,
+                          and not "Wikipedia" in s
+                          and not "wikipedia" in s,
                 [
                     " ".join(re.findall(r"[\w']+", s.strip()))[0:200]
                     for s in soup.get_text().splitlines()
@@ -71,7 +76,7 @@ def create_strings_from_wikipedia(minimum_length, count, lang):
         )
 
         # Remove the last lines that talks about contributing
-        sentences.extend(lines[0 : max([1, len(lines) - 5])])
+        sentences.extend(lines[0: max([1, len(lines) - 5])])
 
     return sentences[0:count]