Repository: roskakori/pygount
Pull request #101 (2023-01-08 12:31:17)

Allow analyzing code from an open file
This implements the feature requested in #100.


Patch:
diff --git a/pygount/analysis.py b/pygount/analysis.py
index 37162ad..3d234a1 100644
--- a/pygount/analysis.py
+++ b/pygount/analysis.py
@@ -12,6 +12,7 @@ import logging
 import os
 import re
 from enum import Enum
+from io import BufferedIOBase, IOBase, RawIOBase, TextIOBase
 from typing import Dict, Generator, List, Optional, Pattern, Sequence, Set, Tuple, Union
 
 import pygments.lexer
@@ -261,10 +262,11 @@ class SourceAnalysis:
         fallback_encoding: str = "cp1252",
         generated_regexes=pygount.common.regexes_from(DEFAULT_GENERATED_PATTERNS_TEXT),
         duplicate_pool: Optional[DuplicatePool] = None,
+        file_handle: Optional[IOBase] = None,
     ) -> "SourceAnalysis":
         """
         Factory method to create a :py:class:`SourceAnalysis` by analyzing
-        the source code in ``source_path``.
+        the source code in ``source_path`` or the open file ``file_handle``.
 
         :param source_path: path to source code to analyze
         :param group: name of a logical group the sourc code belongs to, e.g. a
@@ -276,6 +278,9 @@ class SourceAnalysis:
           if a source code identify is as generated source code for which SLOC should not be counted
         :param duplicate_pool: a :class:`DuplicatePool` where information about possible duplicates is
           collected, or ``None`` if possible duplicates should be counted multiple times.
+        :param file_handle: a file-like object, or ``None`` to read and open the file from
+          ``source_path``. If the file is open in text mode, it must be opened with the correct
+          encoding.
         """
         assert encoding is not None
         assert generated_regexes is not None
@@ -283,30 +288,42 @@ class SourceAnalysis:
         result = None
         lexer = None
         source_code = None
-        source_size = os.path.getsize(source_path)
-        if source_size == 0:
-            _log.info("%s: is empty", source_path)
-            result = SourceAnalysis.from_state(source_path, group, SourceState.empty)
-        elif is_binary_file(source_path):
-            _log.info("%s: is binary", source_path)
-            result = SourceAnalysis.from_state(source_path, group, SourceState.binary)
-        elif not has_lexer(source_path):
-            _log.info("%s: unknown language", source_path)
-            result = SourceAnalysis.from_state(source_path, group, SourceState.unknown)
-        elif duplicate_pool is not None:
+        if file_handle is None:
+            source_size = os.path.getsize(source_path)
+            if source_size == 0:
+                _log.info("%s: is empty", source_path)
+                result = SourceAnalysis.from_state(source_path, group, SourceState.empty)
+            elif is_binary_file(source_path):
+                _log.info("%s: is binary", source_path)
+                result = SourceAnalysis.from_state(source_path, group, SourceState.binary)
+            elif not has_lexer(source_path):
+                _log.info("%s: unknown language", source_path)
+                result = SourceAnalysis.from_state(source_path, group, SourceState.unknown)
+        if duplicate_pool is not None:
             duplicate_path = duplicate_pool.duplicate_path(source_path)
             if duplicate_path is not None:
                 _log.info("%s: is a duplicate of %s", source_path, duplicate_path)
                 result = SourceAnalysis.from_state(source_path, group, SourceState.duplicate, duplicate_path)
         if result is None:
-            if encoding in ("automatic", "chardet"):
-                encoding = encoding_for(source_path, encoding, fallback_encoding)
-            try:
-                with open(source_path, encoding=encoding) as source_file:
-                    source_code = source_file.read()
-            except (LookupError, OSError, UnicodeError) as error:
-                _log.warning("cannot read %s using encoding %s: %s", source_path, encoding, error)
-                result = SourceAnalysis.from_state(source_path, group, SourceState.error, error)
+            if file_handle is None:
+                if encoding in ("automatic", "chardet"):
+                    encoding = encoding_for(source_path, encoding, fallback_encoding)
+                try:
+                    with open(source_path, encoding=encoding) as source_file:
+                        source_code = source_file.read()
+                except (LookupError, OSError, UnicodeError) as error:
+                    _log.warning("cannot read %s using encoding %s: %s", source_path, encoding, error)
+                    result = SourceAnalysis.from_state(source_path, group, SourceState.error, error)
+            elif not isinstance(file_handle, TextIOBase):
+                if encoding in ("automatic", "chardet"):
+                    encoding = encoding_for(source_path, encoding, fallback_encoding, file_handle=file_handle)
+                try:
+                    source_code = file_handle.read().decode(encoding)
+                except (LookupError, OSError, UnicodeError) as error:
+                    _log.warning("cannot read %s using encoding %s: %s", source_path, encoding, error)
+                    result = SourceAnalysis.from_state(source_path, group, SourceState.error, error)
+            else:
+                source_code = file_handle.read()
             if result is None:
                 lexer = guess_lexer(source_path, source_code)
                 assert lexer is not None
@@ -686,7 +703,12 @@ def _line_parts(lexer: pygments.lexer.Lexer, text: str) -> Generator[Set[str], N
         yield line_marks
 
 
-def encoding_for(source_path: str, encoding: str = "automatic", fallback_encoding: Optional[str] = None) -> str:
+def encoding_for(
+    source_path: str,
+    encoding: str = "automatic",
+    fallback_encoding: Optional[str] = None,
+    file_handle: Optional[Union[BufferedIOBase, RawIOBase]] = None,
+) -> str:
     """
     The encoding used by the text file stored in ``source_path``.
 
@@ -706,8 +728,13 @@ def encoding_for(source_path: str, encoding: str = "automatic", fallback_encodin
     assert encoding is not None
 
     if encoding == "automatic":
-        with open(source_path, "rb") as source_file:
-            heading = source_file.read(128)
+        if file_handle is None:
+            with open(source_path, "rb") as source_file:
+                heading = source_file.read(128)
+        else:
+            file_handle.seekable and file_handle.seek(0)
+            heading = file_handle.read(128)
+            file_handle.seekable and file_handle.seek(0)
         result = None
         if len(heading) == 0:
             # File is empty, assume a dummy encoding.
@@ -737,11 +764,17 @@ def encoding_for(source_path: str, encoding: str = "automatic", fallback_encodin
             _detector is not None
         ), 'without chardet installed, encoding="chardet" must be rejected before calling encoding_for()'
         _detector.reset()
-        with open(source_path, "rb") as source_file:
-            for line in source_file.readlines():
-                _detector.feed(line)
-                if _detector.done:
-                    break
+        if file_handle is None:
+            with open(source_path, "rb") as source_file:
+                lines = source_file.readlines()
+        else:
+            file_handle.seekable and file_handle.seek(0)
+            lines = file_handle.readlines()
+            file_handle.seekable and file_handle.seek(0)
+        for line in lines:
+            _detector.feed(line)
+            if _detector.done:
+                break
         result = _detector.result["encoding"]
         if result is None:
             _log.warning(
@@ -759,8 +792,13 @@ def encoding_for(source_path: str, encoding: str = "automatic", fallback_encodin
         else:
             try:
                 # Attempt to read the file as UTF-8.
-                with open(source_path, encoding="utf-8") as source_file:
-                    source_file.read()
+                if file_handle is None:
+                    with open(source_path, encoding="utf-8") as source_file:
+                        source_file.read()
+                else:
+                    file_handle.seekable and file_handle.seek(0)
+                    file_handle.read()
+                    file_handle.seekable and file_handle.seek(0)
                 result = "utf-8"
             except UnicodeDecodeError:
                 # UTF-8 did not work out, use the default as last resort.


Test patch:
diff --git a/tests/test_analysis.py b/tests/test_analysis.py
index d007e82..674a8b1 100644
--- a/tests/test_analysis.py
+++ b/tests/test_analysis.py
@@ -6,6 +6,7 @@ Tests for pygount source code analysis.
 import glob
 import os
 import unittest
+from io import BytesIO, StringIO
 from typing import List, Set
 
 import pytest
@@ -198,6 +199,22 @@ class FileAnalysisTest(TempFolderTest):
         assert source_analysis.state == analysis.SourceState.binary
         assert source_analysis.code_count == 0
 
+    def test_can_analyze_stringio(self):
+        test_path = "imaginary/path/to/file.py"
+        test_code = "from random import randint\n\n# Print a random dice roll\nprint(randint(6))\n"
+        source_analysis = analysis.SourceAnalysis.from_file(test_path, "test", file_handle=StringIO(test_code))
+        assert source_analysis.state == analysis.SourceState.analyzed
+        assert source_analysis.language == "Python"
+        assert source_analysis.code_count == 2
+
+    def test_can_analyze_bytesio(self):
+        test_path = "imaginary/path/to/file.py"
+        test_code = b"from random import randint\n\n# Print a random dice roll\nprint(randint(6))\n"
+        source_analysis = analysis.SourceAnalysis.from_file(test_path, "test", file_handle=BytesIO(test_code))
+        assert source_analysis.state == analysis.SourceState.analyzed
+        assert source_analysis.language == "Python"
+        assert source_analysis.code_count == 2
+
 
 def test_can_repr_source_analysis_from_file():
     source_analysis = analysis.SourceAnalysis("some.py", "Python", "some", 1, 2, 3, 4, analysis.SourceState.analyzed)