llama-index

1
import os
2
from pathlib import Path
3
from typing import List, Optional, Tuple
4

5
import requests
6

7

8
def get_file_content(url: str, path: str) -> Tuple[str, int]:
9
    """Get the content of a file from the GitHub REST API."""
10
    resp = requests.get(url + path)
11
    return resp.text, resp.status_code
12

13

14
def get_file_content_bytes(url: str, path: str) -> Tuple[bytes, int]:
15
    """Get the content of a file from the GitHub REST API."""
16
    resp = requests.get(url + path)
17
    return resp.content, resp.status_code
18

19

20
def get_exports(raw_content: str) -> List:
21
    """Read content of a Python file and returns a list of exported class names.
22

23
    For example:
24
    ```python
25
    from .a import A
26
    from .b import B
27

28
    __all__ = ["A", "B"]
29
    ```
30
    will return `["A", "B"]`.
31

32
    Args:
33
        - raw_content: The content of a Python file as a string.
34

35
    Returns:
36
        A list of exported class names.
37

38
    """
39
    exports = []
40
    for line in raw_content.splitlines():
41
        line = line.strip()
42
        if line.startswith("__all__"):
43
            exports = line.split("=")[1].strip().strip("[").strip("]").split(",")
44
            exports = [export.strip().strip("'").strip('"') for export in exports]
45
    return exports
46

47

48
def rewrite_exports(exports: List[str], dirpath: str) -> None:
49
    """Write the `__all__` variable to the `__init__.py` file in the modules dir.
50

51
    Removes the line that contains `__all__` and appends a new line with the updated
52
    `__all__` variable.
53

54
    Args:
55
        - exports: A list of exported class names.
56

57
    """
58
    init_path = f"{dirpath}/__init__.py"
59
    with open(init_path) as f:
60
        lines = f.readlines()
61
    with open(init_path, "w") as f:
62
        for line in lines:
63
            line = line.strip()
64
            if line.startswith("__all__"):
65
                continue
66
            f.write(line + os.linesep)
67
        f.write(f"__all__ = {list(set(exports))}" + os.linesep)
68

69

70
def initialize_directory(
71
    custom_path: Optional[str] = None, custom_dir: Optional[str] = None
72
) -> Path:
73
    """Initialize directory."""
74
    if custom_path is not None and custom_dir is not None:
75
        raise ValueError(
76
            "You cannot specify both `custom_path` and `custom_dir` at the same time."
77
        )
78

79
    custom_dir = custom_dir or "llamadatasets"
80
    if custom_path is not None:
81
        dirpath = Path(custom_path)
82
    else:
83
        dirpath = Path(__file__).parent / custom_dir
84
    if not os.path.exists(dirpath):
85
        # Create a new directory because it does not exist
86
        os.makedirs(dirpath)
87

88
    return dirpath
89