Add typing

Author: ZipFile

getdents/__init__.py

@@ -1,6 +1,7 @@
 import os
+from typing import Iterator, Tuple
 
-from ._getdents import (  # noqa: ignore=F401
+from ._getdents import (
     DT_BLK,
     DT_CHR,
     DT_DIR,
@@ -14,8 +15,10 @@
     getdents_raw,
 )
 
+DirectoryEntry = Tuple[int, int, str]
 
-def getdents(path, buff_size=32768):
+
+def getdents(path: str, buff_size: int = 32768) -> Iterator[DirectoryEntry]:
     """Get directory entries.
 
     Wrapper around getdents_raw(), simulates ls behaviour: ignores deleted
@@ -48,3 +51,19 @@ def getdents(path, buff_size=32768):
         )
     finally:
         os.close(fd)
+
+
+__all__ = [
+    'DT_BLK',
+    'DT_CHR',
+    'DT_DIR',
+    'DT_FIFO',
+    'DT_LNK',
+    'DT_REG',
+    'DT_SOCK',
+    'DT_UNKNOWN',
+    'MIN_GETDENTS_BUFF_SIZE',
+    'O_GETDENTS',
+    'getdents',
+    'getdents_raw',
+]

getdents/_getdents.pyi

@@ -0,0 +1,14 @@
+from typing import Iterator, Tuple
+
+DT_BLK: int
+DT_CHR: int
+DT_DIR: int
+DT_FIFO: int
+DT_LNK: int
+DT_REG: int
+DT_SOCK: int
+DT_UNKNOWN: int
+MIN_GETDENTS_BUFF_SIZE: int
+O_GETDENTS: int
+
+def getdents_raw(fd: int, buff_size: int) -> Iterator[Tuple[int, int, str]]: ...

getdents/cli.py

@@ -1,11 +1,15 @@
 from argparse import ArgumentParser
-from sys import stderr
+from sys import stderr, stdout
+from typing import List, Optional, Tuple
 
 from . import MIN_GETDENTS_BUFF_SIZE, getdents
-from .formatters import FORMATTERS
+from .formatters import FORMATTERS, Formatter
 
 
-def parse_args(args, prog):
+def parse_args(
+    args: Optional[List[str]],
+    prog: Optional[str],
+) -> Tuple[str, int, Formatter]:
     parser = ArgumentParser(
         prog=prog,
         description='Print directory contents.',
@@ -38,11 +42,11 @@ def parse_args(args, prog):
     return parsed_args.path, buff_size, FORMATTERS[parsed_args.output_format]
 
 
-def main(args=None, prog=None):
+def main(args: Optional[List[str]] = None, prog: Optional[str] = None) -> int:
     path, buff_size, fmt = parse_args(args, prog)
 
     try:
-        fmt(getdents(path, buff_size=buff_size))
+        fmt(getdents(path, buff_size=buff_size), stdout)
     except MemoryError:
         print(
             'Not enough memory to allocate', buff_size, 'bytes of data',

getdents/formatters.py

@@ -2,8 +2,9 @@
 from itertools import chain
 from json import dumps as json_dumps
 from sys import stdout
+from typing import Callable, Dict, Iterable, TextIO
 
-from ._getdents import (
+from . import (
     DT_BLK,
     DT_CHR,
     DT_DIR,
@@ -12,11 +13,13 @@
     DT_REG,
     DT_SOCK,
     DT_UNKNOWN,
+    DirectoryEntry,
 )
 
 
+Formatter = Callable[[Iterable[DirectoryEntry], TextIO], None]
 HEADER = ('inode', 'type', 'name')
-FORMATTERS = {}
+FORMATTERS: Dict[str, Formatter] = {}
 TYPE_NAMES = {
     DT_BLK: 'blk',
     DT_CHR: 'chr',
@@ -29,26 +32,35 @@
 }
 
 
-def formatter(name, registry=FORMATTERS):
-    def deco(fn):
+def formatter(
+    name: str,
+    registry: Dict[str, Formatter] = FORMATTERS,
+) -> Callable[[Formatter], Formatter]:
+    def deco(fn: Formatter) -> Formatter:
         registry[name] = fn
         return fn
     return deco
 
 
 @formatter('plain')
-def format_plain(directory_entries, file=stdout):
+def format_plain(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+) -> None:
     for inode, type, name in directory_entries:
         print(name, file=file)
 
 
 class Echo:
-    def write(self, value):
+    def write(self, value: str) -> str:
         return value
 
 
-@formatter('csv')
-def format_csv(directory_entries, file=stdout, headers=False):
+def _format_csv(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+    headers: bool = False,
+) -> None:
     writer = csv_writer(Echo())
 
     for first in directory_entries:
@@ -62,12 +74,23 @@ def format_csv(directory_entries, file=stdout, headers=False):
             )
 
 
+@formatter('csv')
+def format_csv(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+) -> None:
+    return _format_csv(directory_entries, file, False)
+
+
 @formatter('csv-headers')
-def format_csv_headers(directory_entries, file=stdout):
-    return format_csv(directory_entries, file=file, headers=True)
+def format_csv_headers(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+) -> None:
+    return _format_csv(directory_entries, file, True)
 
 
-def json_encode(inode, type, name):
+def json_encode(inode: int, type: int, name: str) -> str:
     return json_dumps({
         'inode': inode,
         'type': TYPE_NAMES[type],
@@ -76,7 +99,10 @@ def json_encode(inode, type, name):
 
 
 @formatter('json')
-def format_json(directory_entries, file=stdout):
+def format_json(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+) -> None:
     for inode, type, name in directory_entries:
         print(
             '[\n', json_encode(inode, type, name),
@@ -93,6 +119,9 @@ def format_json(directory_entries, file=stdout):
 
 
 @formatter('json-stream')
-def format_json_stream(directory_entries, file=stdout):
+def format_json_stream(
+    directory_entries: Iterable[DirectoryEntry],
+    file: TextIO = stdout,
+) -> None:
     for inode, type, name in directory_entries:
         print(json_encode(inode, type, name), file=file)

getdents/py.typed

@@ -24,6 +24,7 @@
     license='BSD-2-Clause',
     packages=find_packages(exclude=['tests']),
     include_package_data=True,
+    package_data={'getdents': ['py.typed']},
     zip_safe=False,
     extras_require={
         'test': tests_require,

setup.py

@@ -1,4 +1,5 @@
 import re
+from sys import stdout
 
 import pretend
 
@@ -40,7 +41,7 @@ def test_main(monkeypatch):
     directory_entries = pretend.stub()
 
     @pretend.call_recorder
-    def format_test(directory_entries):
+    def format_test(directory_entries, file):
         pass
 
     @pretend.call_recorder
@@ -52,7 +53,7 @@ def getdents(path, buff_size=32768):
 
     assert main(['x', '-o', 'test', '-b', '1024'], 'test') == 0
     assert getdents.calls == [pretend.call('x', buff_size=1024)]
-    assert format_test.calls == [pretend.call(directory_entries)]
+    assert format_test.calls == [pretend.call(directory_entries, stdout)]
 
 
 def test_main_memory_error(monkeypatch):