python-jack-knife 0.5.0__py3-none-any.whl

pjk/sinks/sinks.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import sys
+import yaml
+import subprocess
+import shutil
+from pjk.base import Sink, Source, ParsedToken, Usage
+
+class YamlSink(Sink):
+    @classmethod
+    def usage(cls):
+        usage = Usage(
+            name='yaml',
+            desc='Write all records to a YAML file as multi-doc stream'
+        )
+        usage.def_arg(name='path', usage='Path to output YAML file')
+        return usage
+
+    def __init__(self, ptok: ParsedToken, usage: Usage):
+        super().__init__(ptok, usage)
+        self.path = usage.get_arg('path')
+
+    def process(self) -> None:
+        with open(self.path, 'w') as f:
+            yaml.dump_all(self.input, f, sort_keys=False)
+
+
+class StdoutYamlSink(Sink):
+    # No usage() — not token-based; intended for internal use
+    def __init__(self, input_source: Source, use_pager: bool = True):
+        super().__init__(input_source)
+        self.use_pager = use_pager
+        self.suppress_report = True
+
+    def process(self) -> None:
+        output_stream = sys.stdout
+        pager_proc = None
+
+        if self.use_pager and shutil.which("less"):
+            pager_proc = subprocess.Popen(
+                ["less", "-FRSX"],
+                stdin=subprocess.PIPE,
+                text=True
+            )
+            output_stream = pager_proc.stdin
+
+        try:
+            for record in self.input:
+                try:
+                    yaml.dump(
+                        record,
+                        output_stream,
+                        sort_keys=False,
+                        explicit_start=True,
+                        width=float("inf")
+                    )
+                except BrokenPipeError:
+                    break  # user quit pager
+        except BrokenPipeError:
+            pass
+        finally:
+            if pager_proc:
+                try:
+                    output_stream.close()
+                except BrokenPipeError:
+                    pass
+                pager_proc.wait()

pjk/sinks/stdout.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import sys
+import yaml
+from pjk.base import Sink, Source, ParsedToken, Usage
+from pjk.common import pager_stdout
+
+class StdoutSink(Sink):
+    @classmethod
+    def usage(cls):
+        usage = Usage(
+            name='-',
+            desc='display records in yaml format to stdout through less',
+            component_class=cls
+        )
+        usage.def_param('less', usage='use less to display', valid_values=['true', 'false'], default='true')
+        usage.def_example(["{hello:'world!'}"], "{hello:'world!'}")
+        return usage
+    
+    def __init__(self, ptok: ParsedToken, usage: Usage):
+        super().__init__(ptok, usage)
+
+        # NOTE: self.use_pager is hardcoded for now; override via constructor if needed
+        self.use_pager = True if usage.get_param('less') == None else usage.get_param('less') == 'true'
+
+    def process(self) -> None:
+        # Route all stdout into `less` while dumping
+        try:
+            with pager_stdout(self.use_pager):
+                for record in self.input:
+                    try:
+                        yaml.dump(
+                            record,
+                            sys.stdout,          # now points to less (if enabled)
+                            sort_keys=False,
+                            explicit_start=True,
+                            width=10**9          # effectively no wrap without using a float
+                        )
+                    except BrokenPipeError:
+                        break  # user quit pager
+        except BrokenPipeError:
+            # Swallow if pager closed early
+            pass

pjk/sinks/tsv_sink.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+from pjk.sinks.csv_sink import CSVSink
+from pjk.base import Source, ParsedToken, Usage
+
+class TSVSink(CSVSink):
+    is_format = True
+
+    @classmethod
+    def usage(cls):
+        usage = Usage(
+            name='tsv',
+            desc='Write records to a .tsv file (tab-separated values)',
+            component_class=cls
+        )
+        usage.def_arg(name='path', usage='Path prefix (no extension)')
+        return usage
+
+    def __init__(self, ptok: ParsedToken, usage: Usage):
+        path_no_ext = usage.get_arg('path')
+        super().__init__(path_no_ext, delimiter="\t", ext='tsv')

pjk/sinks/user_sink_factory.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import importlib.util
+from typing import Optional
+from pjk.base import Source, Sink, UsageError, ParsedToken
+
+class UserSinkFactory:
+    @staticmethod
+    def create_from_path(ptok: ParsedToken) -> Optional[Sink]:
+        script_path = ptok.pre_colon
+        try:
+            spec = importlib.util.spec_from_file_location("user_sink", script_path)
+            if spec is None or spec.loader is None:
+                raise UsageError(f"Could not load Python file: {script_path}")
+
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+        except Exception as e:
+            raise UsageError(f"Failed to import {script_path}: {e}")
+
+        for value in vars(module).values():
+            if (
+                isinstance(value, type)
+                and issubclass(value, Sink)
+                and value is not Sink
+                and value.__module__ == module.__name__
+            ):
+                usage = value.usage()
+                usage.bind(ptok)
+
+                return value(ptok, usage)
+
+        return None
+
+    @classmethod
+    def create(cls, ptok: ParsedToken) -> Sink:
+        if ptok.pre_colon.endswith('.py'):
+            sink = cls.create_from_path(ptok)
+            if sink:
+                return sink
+
+        

pjk/sources/csv_source.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import sys
+import csv
+from pjk.base import Source, NoBindUsage
+from pjk.sources.format_usage import FormatUsage
+from pjk.sources.lazy_file import LazyFile
+
+csv.field_size_limit(sys.maxsize)
+
+class CSVSource(Source):
+    is_format = True
+    @classmethod
+    def usage(cls):
+        return FormatUsage('csv', component_class=cls)
+
+    def __init__(self, lazy_file: LazyFile, delimiter: str = ","):
+        self.lazy_file = lazy_file
+        self.delimiter = delimiter
+        self.num_recs = 0
+
+    def __iter__(self):
+        with self.lazy_file.open() as f:
+            reader = csv.DictReader(f, delimiter=self.delimiter)
+            for row in reader:
+                self.num_recs += 1
+                yield row

pjk/sources/dir_source.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import os
+from typing import Any
+from queue import Queue, Empty
+from pjk.base import Source, ParsedToken
+from pjk.sources.lazy_file_local import LazyFileLocal
+from pjk.log import logger
+
+class DirSource(Source):
+    def __init__(self, source_queue: Queue, in_source: Source = None):
+        self.source_queue = source_queue
+        self.current = in_source
+
+    def __iter__(self):
+        while True:
+            if self.current is None:
+                try:
+                    self.current = self.source_queue.get_nowait()
+                    logger.debug(f'next source={self.current}')
+                except Empty:
+                    return  # end of all sources
+
+            try:
+                for record in self.current:
+                    yield record
+            finally:
+                self.current = None  # move to next source after exhaustion
+
+    def deep_copy(self):
+        if self.source_queue.qsize() <= 1:
+            return None  # leave remaining files to original
+        try:
+            next_source = self.source_queue.get_nowait()
+            logger.debug(f'deep_copy next_source={next_source}')
+        except Empty:
+            return None
+
+        return DirSource(self.source_queue, next_source)
+
+    @classmethod
+    def create(cls, ptok: ParsedToken, get_format_class_gz: Any):
+        params = ptok.get_params()
+        override = params.get('format', None)
+        path = ptok.all_but_params
+
+        files = [
+            os.path.join(path, f)
+            for f in os.listdir(path)
+            if os.path.isfile(os.path.join(path, f))
+        ]
+
+        source_queue = Queue()
+        for file in files:
+            file_token = file if not override else f"{file}@format={override}"
+            file_ptok = ParsedToken(file_token)
+
+            format_class, is_gz = get_format_class_gz(file_ptok)
+            if format_class:
+                lazy_file = LazyFileLocal(file, is_gz)
+                source_queue.put(format_class(lazy_file))
+            else:
+                raise RuntimeError(f"No format for file: {file}")
+
+        if source_queue.empty():
+            return None
+
+        return DirSource(source_queue)

pjk/sources/factory.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import os
+import queue
+from pjk.base import Source, ParsedToken
+from pjk.common import ComponentFactory
+from pjk.sources.json_source import JsonSource
+from pjk.sources.csv_source import CSVSource
+from pjk.sources.sql_source import SQLSource
+from pjk.sources.tsv_source import TSVSource
+from pjk.sources.s3_source import S3Source
+from pjk.sources.inline_source import InlineSource
+from pjk.sources.dir_source import DirSource
+from pjk.sources.user_source_factory import UserSourceFactory
+from pjk.sources.lazy_file import LazyFile
+from pjk.sources.lazy_file_local import LazyFileLocal
+from pjk.sources.parquet_source import ParquetSource
+
+COMPONENTS = {
+        'inline': InlineSource,
+        'json': JsonSource,
+        'csv': CSVSource,
+        'tsv': TSVSource,
+        'sql': SQLSource,
+        'parquet': ParquetSource,
+    }
+
+class SourceFactory(ComponentFactory):
+    def __init__(self):
+        super().__init__(COMPONENTS, 'source')
+    
+
+    def get_format_class_gz(self, ptok: ParsedToken):
+        params = ptok.get_params()
+        override = params.get('format', None) # e.g. json or json.gz
+
+        lookup = None
+
+        is_gz = ptok.all_but_params.endswith('gz')
+        if override:
+            if override.endswith('.gz'):
+                is_gz = True
+                override = override.removesuffix('.gz')
+            lookup = override
+
+        else: # e.g. foo.json or foo.json.gz
+            path = ptok.all_but_params
+            if path.endswith('.gz'):
+                is_gz = True
+                path = path.removesuffix('.gz')
+
+            path, ext = os.path.splitext(path) # e.g path=foo.json
+            lookup = ext.removeprefix('.')
+            
+        format_class = self.components.get(lookup, None)
+        if not format_class:
+            return None, None
+        
+        # make sure
+        if not format_class.is_format:
+            return None, None # raise ?
+
+        return format_class, is_gz
+
+    def create(self, token: str) -> Source:
+        token = token.strip()
+
+        if InlineSource.is_inline(token):
+            return InlineSource(token)
+        
+        ptok = ParsedToken(token)
+        
+        if ptok.pre_colon.endswith('.py'):
+            source = UserSourceFactory.create(ptok)
+            if source:
+                return source
+
+        source_cls = self.components.get(ptok.pre_colon)
+        if source_cls:
+            usage = source_cls.usage()
+            usage.bind(ptok)
+        
+            source = source_cls(ptok, usage)
+            return source
+
+        if ptok.all_but_params.startswith('s3'):
+            return S3Source.create(ptok, get_format_class_gz=self.get_format_class_gz)
+
+        if os.path.isdir(ptok.all_but_params):
+            return DirSource.create(ptok, get_format_class_gz=self.get_format_class_gz)
+
+        # individual file
+        if os.path.isfile(ptok.all_but_params):
+            source_class, is_gz = self.get_format_class_gz(ptok)
+            if source_class:
+                lazy_file = LazyFileLocal(ptok.all_but_params, is_gz)
+                return source_class(lazy_file)
+     
+        return None

pjk/sources/format_usage.py

@@ -0,0 +1,11 @@
+from pjk.base import Source, NoBindUsage
+
+class FormatUsage(NoBindUsage):
+    def __init__(self, name: str, component_class: type, desc_override: str = None):
+        desc = f'{name} source for s3 and local files/directories.' if desc_override == None else desc_override
+        super().__init__(name, desc, component_class)
+
+        self.def_syntax("") # no syntax for these
+        self.def_example(expr_tokens=[f"myfile.{name}", "-"], expect=None)
+        self.def_example(expr_tokens=["mydir", "-"], expect=None)
+        self.def_example(expr_tokens=["s3://mybucket/path/to/files", "-"], expect=None)

pjk/sources/inline_source.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+# djk/sources/inline_source.py
+
+import hjson
+from hjson import HjsonDecodeError
+from typing import Optional
+from collections import OrderedDict
+from pjk.base import Source, TokenError, Usage
+
+def to_builtin(obj):
+    """Recursively convert OrderedDicts to dicts and lists."""
+    if isinstance(obj, OrderedDict):
+        return {k: to_builtin(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [to_builtin(v) for v in obj]
+    else:
+        return obj
+
+class InlineSource(Source):
+    @classmethod
+    def usage(cls):
+        usage = Usage(
+            name='inline',
+            desc="simplified json lines format (uses hjson)",
+            component_class=cls
+        )
+        usage.def_syntax('')
+        usage.def_example(expr_tokens=["{hello: 'world!'}"], expect="{hello: 'world!'}")
+        usage.def_example(expr_tokens=["[{id:1, dir:'up'},{id:2, dir:'down'}]"], expect="[{id:1, dir:'up'}, {id:2, dir:'down'}]")
+        return usage
+
+    def __init__(self, inline_expr):
+        self.num_recs = 0
+        try:
+            obj = hjson.loads(inline_expr)
+        except HjsonDecodeError:
+            raise TokenError('incorrect hjson line syntax')
+
+        if isinstance(obj, dict):
+            self.records = [obj]
+        elif isinstance(obj, list):
+            self.records = obj
+        else:
+            raise TokenError(f'"{inline_expr}"')
+
+    def __iter__(self):
+        for raw in self.records:
+            yield to_builtin(raw)
+
+    @classmethod
+    def is_inline(cls, token):
+        if len(token) < 2:
+            return False
+        return (token[0], token[-1]) in {('{', '}'), ('[', ']')}

pjk/sources/json_source.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import json
+from pjk.base import Source, NoBindUsage
+from pjk.sources.lazy_file import LazyFile
+from pjk.sources.format_usage import FormatUsage
+from pjk.log import logger
+
+class JsonSource(Source):
+    is_format = True
+
+    @classmethod
+    def usage(cls):
+        return FormatUsage('json', component_class=cls)
+
+    def __init__(self, lazy_file: LazyFile):
+        self.lazy_file = lazy_file
+        self.num_recs = 0
+
+    def __iter__(self):
+        with self.lazy_file.open() as f:
+            for line in f:
+                self.num_recs += 1
+                try:
+                    yield json.loads(line)
+                except json.JSONDecodeError as e:
+                    print('json decode error, see ~/.pjk/logs')
+                    snippet = line.strip()
+                    if len(snippet) > 200:
+                        snippet = snippet[:200] + "…"
+                    logger.warning(
+                        f"Skipping invalid JSON at line {self.num_recs} "
+                        f"in {self.lazy_file.path}: {e} | data: {snippet}"
+                    )

pjk/sources/lazy_file.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+from abc import ABC, abstractmethod
+from typing import IO
+
+class LazyFile(ABC):
+    @abstractmethod
+    def open(self) -> IO[str]:
+        """Open and return a text-mode file-like object."""
+        pass
+
+    @abstractmethod
+    def name(self) -> str:
+        """Return a descriptive identifier (e.g. path or URI)."""
+        pass

pjk/sources/lazy_file_local.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import gzip
+import io
+from typing import IO
+from pjk.sources.lazy_file import LazyFile
+
+class LazyFileLocal(LazyFile):
+    def __init__(self, path: str, is_gz: bool = False):
+        self.path = path
+        self.is_gz = is_gz
+
+    def open(self) -> IO[str]:
+        raw = open(self.path, "rb")
+        if self.path.endswith(".gz") or self.is_gz:
+            return io.TextIOWrapper(gzip.GzipFile(fileobj=raw))
+        else:
+            return io.TextIOWrapper(raw)
+
+    def name(self) -> str:
+        return self.path

pjk/sources/lazy_file_s3.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+import io
+import gzip
+from typing import IO
+from pjk.sources.lazy_file import LazyFile
+
+class LazyFileS3(LazyFile):
+    def __init__(self, bucket: str, key: str, is_gz: bool):
+        import boto3 # lazy import
+        self.s3 = boto3.client('s3') # for each thread
+        self.bucket = bucket
+        self.key = key
+        self.is_gz = is_gz
+
+    def open(self, binary=False) -> IO[str]:
+        obj = self.s3.get_object(Bucket=self.bucket, Key=self.key)
+        raw_body = obj['Body'].read()
+        if self.is_gz:
+            return io.TextIOWrapper(gzip.GzipFile(fileobj=io.BytesIO(raw_body)))
+        elif binary:
+            return io.BytesIO(raw_body)
+        else:
+            return io.StringIO(raw_body.decode("utf-8"))
+
+    def name(self) -> str:
+        return f"s3://{self.bucket}/{self.key}"

pjk/sources/parquet_source.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+from pjk.base import Source
+from pjk.sources.lazy_file import LazyFile
+from pjk.sources.format_usage import FormatUsage
+
+class ParquetSource(Source):
+    is_format = True  # enables format-based routing
+    @classmethod
+    def usage(cls):
+        return FormatUsage('parquet', component_class=cls)
+
+    def __init__(self, lazy_file: LazyFile):
+        self.lazy_file = lazy_file
+        self.num_recs = 0
+
+    def __iter__(self):
+        import pyarrow.parquet as pq # lazy import
+        with self.lazy_file.open(binary=True) as f:
+            table = pq.read_table(f)
+            batch = table.to_pydict()
+
+            if not batch:
+                return  # no columns = no rows
+
+            num_rows = len(next(iter(batch.values())))
+
+            for i in range(num_rows):
+                record = {col: batch[col][i] for col in batch}
+                self.num_recs += 1
+                yield record