Wrote text ingestors, still incomplete tests (#37)

* Wrote text ingestors, still incomplete tests

* Added tests for text

* removed print

* Resolved conflicts

Author: Ansh5461

querent/ingestors/ingestor_errors.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class IngestorError(Enum):
+    EOF = "End of File"
+    ETCD = "ETCD Error"
+    NETWORK = "Network Error"
+    TIMEOUT = "Timeout"
+    UNKNOWN = "Unknown Error"

querent/ingestors/ingestor_manager.py

@@ -1,17 +1,24 @@
+"""
+Ingestor manager, for managing all the factories with backend
+"""
 from typing import Optional
 from querent.config.ingestor_config import IngestorBackend
 from querent.ingestors.base_ingestor import BaseIngestor
 from querent.ingestors.ingestor_factory import IngestorFactory, UnsupportedIngestor
 from querent.ingestors.pdfs.pdf_ingestor_v1 import PdfIngestorFactory
+from querent.ingestors.texts.text_ingestor import TextIngestorFactory
 from querent.ingestors.audio.audio_ingestors import AudioIngestorFactory
 from querent.ingestors.json.json_ingestor import JsonIngestorFactory
 from querent.ingestors.images.image_ingestor import ImageIngestorFactory
 
 
 class IngestorFactoryManager:
+    """Factory manager"""
+
     def __init__(self):
         self.ingestor_factories = {
             IngestorBackend.PDF.value: PdfIngestorFactory(),
+            IngestorBackend.TEXT.value: TextIngestorFactory(),
             IngestorBackend.MP3.value: AudioIngestorFactory(),
             IngestorBackend.WAV.value: AudioIngestorFactory(),
             IngestorBackend.JSON.value: JsonIngestorFactory(),
@@ -22,14 +29,17 @@ def __init__(self):
         }
 
     async def get_factory(self, file_extension: str) -> IngestorFactory:
+        """get_factory to match factory based on file extension"""
         return self.ingestor_factories.get(
             file_extension.lower(), UnsupportedIngestor("Unsupported file extension")
         )
 
     async def get_ingestor(self, file_extension: str) -> Optional[BaseIngestor]:
+        """get_ingestor to get factory for that extension"""
         factory = self.get_factory(file_extension)
         return factory.create(file_extension)
 
     async def supports(self, file_extension: str) -> bool:
+        """check if extension supports factory"""
         factory = self.get_factory(file_extension)
         return factory.supports(file_extension)

querent/ingestors/texts/text_ingestor.py

@@ -0,0 +1,76 @@
+from typing import List, AsyncGenerator
+from querent.common.types.collected_bytes import CollectedBytes
+from querent.ingestors.base_ingestor import BaseIngestor
+from querent.ingestors.ingestor_factory import IngestorFactory
+from querent.processors.async_processor import AsyncProcessor
+from querent.config.ingestor_config import IngestorBackend
+
+
+class TextIngestorFactory(IngestorFactory):
+    SUPPORTED_EXTENSIONS = {"txt"}
+
+    async def supports(self, file_extension: str) -> bool:
+        return file_extension.lower() in self.SUPPORTED_EXTENSIONS
+
+    async def create(
+        self, file_extension: str, processors: List[AsyncProcessor]
+    ) -> BaseIngestor:
+        if not self.supports(file_extension):
+            return None
+
+        return TextIngestor(processors)
+
+
+class TextIngestor(BaseIngestor):
+    def __init__(self, processors: List[AsyncProcessor]):
+        super().__init__(IngestorBackend.TEXT)
+        self.processors = processors
+
+    async def ingest(
+        self, poll_function: AsyncGenerator[CollectedBytes, None]
+    ) -> AsyncGenerator[List[str], None]:
+        try:
+            collected_bytes = b""
+            current_file = None
+
+            async for chunk_bytes in poll_function:
+                if chunk_bytes.is_error():
+                    continue
+
+                if chunk_bytes.file != current_file:
+                    if current_file:
+                        text = await self.extract_and_process_text(
+                            CollectedBytes(file=current_file, data=collected_bytes)
+                        )
+                        yield text
+
+                    collected_bytes = b""
+                    current_file = chunk_bytes.file
+
+                collected_bytes += chunk_bytes.data
+
+            if current_file:
+                text = await self.extract_and_process_text(
+                    CollectedBytes(file=current_file, data=collected_bytes)
+                )
+                yield text
+
+        except Exception as e:
+            print(e)
+            yield []
+
+    async def extract_and_process_text(
+        self, collected_bytes: CollectedBytes
+    ) -> List[str]:
+        text = await self.extract_text_from_file(collected_bytes)
+        return await self.process_data(text=text)
+
+    async def extract_text_from_file(self, collected_bytes: CollectedBytes) -> str:
+        text = collected_bytes.data.decode("utf-8")
+        return text
+
+    async def process_data(self, text: str) -> List[str]:
+        processed_data = text
+        for processor in self.processors:
+            processed_data = await processor.process(processed_data)
+        return processed_data

tests/data/text/asyncgenerator.txt

@@ -0,0 +1,12 @@
+Asynchronous generator functions are part of Python version 3.6, they were introduced by PEP-525. Asynchronous generator 
+functions are much like regular asynchronous functions except that they contain the yield keyword in the function body.
+Which in turn, makes them much like regular generators, except for that you can use the await keyword in there as well.
+
+When calling an asynchronous generator function, the result that is returned is an asynchronous generator object. In 
+contrast to calling regular asynchronous functions which return a coroutine object.
+Since the asynchronous generator is, no surprise, asynchronous you are allowed to use the await keyword inside the 
+asynchronous generator.
+
+You can use this, for example, to send out HTTP requests in the asynchronous generator and yielding the response.
+
+Besides asynchronous iterables you can use asynchronous generators with the async for-loop as well.

tests/test_audio_ingestor.py

@@ -1,10 +1,12 @@
-import asyncio
+"""Test cases for audio ingestors"""
 from pathlib import Path
+import pytest
+import asyncio
+
 from querent.collectors.fs.fs_collector import FSCollectorFactory
 from querent.config.collector_config import FSCollectorConfig
 from querent.common.uri import Uri
 from querent.ingestors.ingestor_manager import IngestorFactoryManager
-import pytest
 
 
 @pytest.mark.asyncio
@@ -30,7 +32,7 @@ async def poll_and_print():
             if len(ingested) == 0:
                 counter += 1
 
-        assert counter == 0
+    assert counter == 0
 
     await poll_and_print()
 

tests/test_pdf_ingestor.py

@@ -32,7 +32,7 @@ async def poll_and_print():
             assert ingested is not None
             if ingested is not "" or ingested is not None:
                 counter += 1
-        assert counter == 19 # 19 pages in the PDF
+        assert counter == 19  # 19 pages in the PDF
 
     await poll_and_print()  # Notice the use of await here
 

tests/test_text_ingestor.py

@@ -0,0 +1,39 @@
+import asyncio
+from pathlib import Path
+from querent.collectors.fs.fs_collector import FSCollectorFactory
+from querent.config.collector_config import FSCollectorConfig
+from querent.common.uri import Uri
+from querent.ingestors.ingestor_manager import IngestorFactoryManager
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_collect_and_ingest_txt():
+    # Set up the collector
+    collector_factory = FSCollectorFactory()
+    uri = Uri("file://" + str(Path("./tests/data/text/").resolve()))
+    config = FSCollectorConfig(root_path=uri.path)
+    collector = collector_factory.resolve(uri, config)
+
+    # Set up the ingestor
+    ingestor_factory_manager = IngestorFactoryManager()
+    ingestor_factory = await ingestor_factory_manager.get_factory("txt")
+    ingestor = await ingestor_factory.create("txt", [])
+
+    # Collect and ingest the PDF
+    ingested_call = ingestor.ingest(collector.poll())
+    counter = 0
+
+    async def poll_and_print():
+        counter = 0
+        async for ingested in ingested_call:
+            assert ingested is not None
+            if len(ingested) == 0:
+                counter += 1
+        assert counter == 0
+
+    await poll_and_print()  # Notice the use of await here
+
+
+if __name__ == "__main__":
+    asyncio.run(test_collect_and_ingest_txt())