write_po: refactor into generate_po

babel/messages/pofile.py

@@ -529,45 +529,72 @@ def write_po(
                              updating the catalog
     :param include_lineno: include line number in the location comment
     """
+
+    sort_by = None
+    if sort_output:
+        sort_by = "message"
+    elif sort_by_file:
+        sort_by = "location"
+
+    for line in generate_po(
+        catalog,
+        ignore_obsolete=ignore_obsolete,
+        include_lineno=include_lineno,
+        include_previous=include_previous,
+        no_location=no_location,
+        omit_header=omit_header,
+        sort_by=sort_by,
+        width=width,
+    ):
+        if isinstance(line, str):
+            line = line.encode(catalog.charset, 'backslashreplace')
+        fileobj.write(line)
+
+
+def generate_po(
+    catalog: Catalog,
+    *,
+    ignore_obsolete: bool = False,
+    include_lineno: bool = True,
+    include_previous: bool = False,
+    no_location: bool = False,
+    omit_header: bool = False,
+    sort_by: Literal["message", "location"] | None,
+    width: int = 76,
+) -> Iterable[str]:
+    r"""Yield text strings representing a ``gettext`` PO (portable object) file.
+
+    See `write_po()` for a more detailed description.
+    """
+    # xgettext always wraps comments even if --no-wrap is passed;
+    # provide the same behaviour
+    comment_width = width if width and width > 0 else 76
+
     def _normalize(key, prefix=''):
         return normalize(key, prefix=prefix, width=width)
 
-    def _write(text):
-        if isinstance(text, str):
-            text = text.encode(catalog.charset, 'backslashreplace')
-        fileobj.write(text)
-
-    def _write_comment(comment, prefix=''):
-        # xgettext always wraps comments even if --no-wrap is passed;
-        # provide the same behaviour
-        _width = width if width and width > 0 else 76
-        for line in wraptext(comment, _width):
-            _write(f"#{prefix} {line.strip()}\n")
+    def _format_comment(comment, prefix=''):
+        for line in wraptext(comment, comment_width):
+            yield f"#{prefix} {line.strip()}\n"
 
-    def _write_message(message, prefix=''):
+    def _format_message(message, prefix=''):
         if isinstance(message.id, (list, tuple)):
             if message.context:
-                _write(f"{prefix}msgctxt {_normalize(message.context, prefix)}\n")
-            _write(f"{prefix}msgid {_normalize(message.id[0], prefix)}\n")
-            _write(f"{prefix}msgid_plural {_normalize(message.id[1], prefix)}\n")
+                yield f"{prefix}msgctxt {_normalize(message.context, prefix)}\n"
+            yield f"{prefix}msgid {_normalize(message.id[0], prefix)}\n"
+            yield f"{prefix}msgid_plural {_normalize(message.id[1], prefix)}\n"
 
             for idx in range(catalog.num_plurals):
                 try:
                     string = message.string[idx]
                 except IndexError:
                     string = ''
-                _write(f"{prefix}msgstr[{idx:d}] {_normalize(string, prefix)}\n")
+                yield f"{prefix}msgstr[{idx:d}] {_normalize(string, prefix)}\n"
         else:
             if message.context:
-                _write(f"{prefix}msgctxt {_normalize(message.context, prefix)}\n")
-            _write(f"{prefix}msgid {_normalize(message.id, prefix)}\n")
-            _write(f"{prefix}msgstr {_normalize(message.string or '', prefix)}\n")
-
-    sort_by = None
-    if sort_output:
-        sort_by = "message"
-    elif sort_by_file:
-        sort_by = "location"
+                yield f"{prefix}msgctxt {_normalize(message.context, prefix)}\n"
+            yield f"{prefix}msgid {_normalize(message.id, prefix)}\n"
+            yield f"{prefix}msgstr {_normalize(message.string or '', prefix)}\n"
 
     for message in _sort_messages(catalog, sort_by=sort_by):
         if not message.id:  # This is the header "message"
@@ -580,12 +607,12 @@ def _write_message(message, prefix=''):
                     lines += wraptext(line, width=width,
                                       subsequent_indent='# ')
                 comment_header = '\n'.join(lines)
-            _write(f"{comment_header}\n")
+            yield f"{comment_header}\n"
 
         for comment in message.user_comments:
-            _write_comment(comment)
+            yield from _format_comment(comment)
         for comment in message.auto_comments:
-            _write_comment(comment, prefix='.')
+            yield from _format_comment(comment, prefix='.')
 
         if not no_location:
             locs = []
@@ -606,35 +633,34 @@ def _write_message(message, prefix=''):
                     location = f"{location}:{lineno:d}"
                 if location not in locs:
                     locs.append(location)
-            _write_comment(' '.join(locs), prefix=':')
+            yield from _format_comment(' '.join(locs), prefix=':')
         if message.flags:
-            _write(f"#{', '.join(['', *sorted(message.flags)])}\n")
+            yield f"#{', '.join(['', *sorted(message.flags)])}\n"
 
         if message.previous_id and include_previous:
-            _write_comment(
+            yield from _format_comment(
                 f'msgid {_normalize(message.previous_id[0])}',
                 prefix='|',
             )
             if len(message.previous_id) > 1:
-                _write_comment('msgid_plural %s' % _normalize(
-                    message.previous_id[1],
-                ), prefix='|')
+                norm_previous_id = _normalize(message.previous_id[1])
+                yield from _format_comment(f'msgid_plural {norm_previous_id}', prefix='|')
 
-        _write_message(message)
-        _write('\n')
+        yield from _format_message(message)
+        yield '\n'
 
     if not ignore_obsolete:
         for message in _sort_messages(
             catalog.obsolete.values(),
             sort_by=sort_by,
         ):
             for comment in message.user_comments:
-                _write_comment(comment)
-            _write_message(message, prefix='#~ ')
-            _write('\n')
+                yield from _format_comment(comment)
+            yield from _format_message(message, prefix='#~ ')
+            yield '\n'
 
 
-def _sort_messages(messages: Iterable[Message], sort_by: Literal["message", "location"]) -> list[Message]:
+def _sort_messages(messages: Iterable[Message], sort_by: Literal["message", "location"] | None) -> list[Message]:
     """
     Sort the given message iterable by the given criteria.