python - 如何在 mypy 中定义隐含导入？

Question

我正在尝试对从 Databricks 导出的笔记本进行类型检查。笔记本是普通的 *.py 文件，带有特殊的注释格式以指示单元格的开始和结束位置。 mypy 没有理由不能对这些文件进行类型检查，除了一些缺失的名称:

• Spark
• sc
• dbutils
• 显示
• 显示HTML

我知道 python 命令将在将您转储到交互模式之前运行由 PYTHONSTARTUP 环境变量指定的文件。这就是这些名称的定义方式。

mypy 中是否有一个钩子(Hook)可以让您在代码之外定义这样的名称？

Answer 1

这是我想出的答案。它很脏，但它有效。我想要一个更好的答案，但在那之前，这是有效的。

策略是使用 shell 脚本将“PYTHONSTARTUP”文件添加到每个笔记本中，然后在最终输出中减去行号。

类型检查.sh:

#!/bin/bash

TARGET=$1

# Define the contents of "PYTHONSTARTUP" file inline. This just
# makes it easier to copy & paste this script elsewhere. You could also 
# make it a separate *.py file.
PRELUDE="$(cat <<EOF
import typing
import pyspark.SparkContext
import pyspark.sql.SparkSession

spark = None  # type: pyspark.sql.SparkSession
sc = None  # type: pyspark.SparkContext

def display(expr):
    pass

def displayHTML(expr):
    pass

class dbutils:
    class fs:
        def help(): pass
        def cp(from_: str, to: str, recurse: bool = False) -> bool: pass
        def head(file: str, maxBytes: int) -> str: pass
        def ls(dir: str) -> typing.List[str]: pass
        def mkdirs(dir: str) -> bool: pass
        def put(file: str, contents: str, overwrite: bool = False) -> bool: pass
        def rm(dir: str, recurse: bool) -> bool: pass
        def mount(source: str, mountPoint: str, encryptionType: str = "", owner: str = "", extraConfigs: typing.Map[str, str] = {}) -> bool: pass
        def mounts() -> typing.List[str]: pass
        def refreshMounts() -> bool: pass
        def unmount(mountPoint: str) -> bool: pass
    class notebook:
        def exit(value: str): pass
        def run(path: str, timeout: int, arguments: typing.Map[str, str]) -> str: pass
    class widgets:
        def combobox(name: str, defaultValue: str, choices: typing.List[str], label: str = ""): pass
        def dropdown(name: str, defaultValue: str, choices: typing.List[str], label: str = ""): pass
        def get(name: str) -> str: pass
        def multiselect(name: str, defaultValue: str, choices: typing.List[str], label: str = ""): pass
        def remove(name: str): pass
        def removeAll(): pass
        def text(name: str, defaultValue: str, label: str = ""): pass

def getArgument(name: str) -> str: pass
EOF
)"

# Remember the length of $PRELUDE so that we can subtract the line number
LEN="$(echo "$PRELUDE" | wc -l | awk '{ print $1 }')"

for file in $(find $TARGET -name '*.py'); do
  # run mypy for the two files concatenated together (with a blank line 
  # for good measure)
  OUTPUT=$(mypy -c "$(cat <<EOF
$PRELUDE

$(cat $file)
EOF
)")
  # awk: Take only output where the line number is after the PRELUDE. Also, fix the file name and line number
  FILE_OUTPUT="$(echo "$OUTPUT" | awk -F: '$2 > '$LEN' { line=($2-'$LEN')-1; $1=""; $2=""; print "'$file':" line ":" $0 }')"

  # Remove blank lines from output before printing
  if [[ $(echo "$FILE_OUTPUT" | sed '/^$/d' | wc -l) -gt 0 ]]; then
    echo "$FILE_OUTPUT"
  fi

  # Keep track of all output, so we can decide the exit code
  ALL_OUTPUT+="$FILE_OUTPUT"
done

# propagate errors to the exit code, but ignore errors in the prelude. This 
# makes it easier to use in a CI pipeline.
if [[ $(echo "$ALL_OUTPUT" | wc -l) -gt 1 ]]; then
  exit 1
else
  exit 0
fi

用法:

./typecheck.sh notebooks/